discretizer.py 文件源码

def _discretize_by_kmeans(col, num_bins, random_state):
    nan_idx = col[col.isnull()].index
    kmeans = KMeans(n_clusters=num_bins, random_state=random_state)
    kmeans = kmeans.fit(col.dropna().values.T.reshape(-1, 1))
    group = kmeans.labels_
    if col.isnull().sum() > 0:
        group = group.astype(float)
        for idx in nan_idx:
            group = np.insert(group,idx,np.nan)
    return pd.Series(group)