如何在scikit-learn层次聚类中设置每个簇的目标样本数？

问题背景 #

我有一组5400个带X、Y坐标的点存储在DataFrame中，希望通过聚类让每个最终簇的样本数在170-230之间（目标200）。当前使用scikit-learn的AgglomerativeClustering层次聚类，但该算法仅支持通过n_clusters指定簇数量或distance_threshold指定最大聚类距离，无法直接设置簇的目标样本数。

现有尝试的代码及结果：

np.random.seed(0)
n_points = 5400
df = pd.DataFrame({
    'X': np.random.randn(n_points),
    'Y': np.random.randn(n_points),
})

model = AgglomerativeClustering(
    n_clusters=27,
    distance_threshold=None,
    linkage="ward"
)

clustered = model.fit_predict(df)
df['clustered'] = clustered

centroids = df.groupby('clustered').agg({'X': 'mean', 'Y': 'mean'}).reset_index()
centroids_df = pd.DataFrame({
    'X': centroids['X'],
    'Y': centroids['Y'],
    'Clustered': centroids['clustered']
})

plt.scatter(df['X'], df['Y'], c=clustered, s=1)
plt.scatter(centroids_df['X'], centroids_df['Y'], c='black', s=200)
plt.show()

聚类结果图：

解决方案 #

一、自定义调整层次聚类结果 #

AgglomerativeClustering本身不支持按样本数约束，需要在初始聚类后手动调整簇大小：

步骤说明

• 统计初始簇样本数：计算每个簇的点数，筛选出过大（>230）和过小（<170）的簇
• 拆分过大簇：对超标的大簇，用K-Means再次聚类，拆分后的子簇数量按簇大小//200计算，确保子簇落在目标范围内
• 合并过小簇：对样本数不足的簇，计算簇间距离，将距离最近的小簇与相邻簇合并，直到所有簇符合要求

代码实现

import numpy as np
import pandas as pd
from sklearn.cluster import AgglomerativeClustering, KMeans
from sklearn.metrics.pairwise import euclidean_distances
import matplotlib.pyplot as plt

np.random.seed(0)
n_points = 5400
df = pd.DataFrame({
    'X': np.random.randn(n_points),
    'Y': np.random.randn(n_points),
})

# 初始层次聚类
model = AgglomerativeClustering(n_clusters=27, linkage="ward")
df['clustered'] = model.fit_predict(df)

# 定义目标范围
MIN_SIZE = 170
MAX_SIZE = 230
TARGET_SIZE = 200

# 拆分过大的簇
def split_large_clusters(df, min_size, max_size, target_size):
    cluster_counts = df['clustered'].value_counts()
    new_cluster_id = df['clustered'].max() + 1
    for cluster_id, count in cluster_counts.items():
        if count > max_size:
            # 计算需要拆分的簇数量
            n_split = count // target_size
            if count % target_size > min_size:
                n_split += 1
            # 提取当前簇的点
            cluster_points = df[df['clustered'] == cluster_id][['X', 'Y']]
            # K-Means拆分
            kmeans = KMeans(n_clusters=n_split, random_state=0)
            sub_clusters = kmeans.fit_predict(cluster_points)
            # 更新簇ID
            df.loc[df['clustered'] == cluster_id, 'clustered'] = sub_clusters + new_cluster_id
            new_cluster_id += n_split
    return df

# 合并过小的簇
def merge_small_clusters(df, min_size):
    while True:
        cluster_counts = df['clustered'].value_counts()
        small_clusters = cluster_counts[cluster_counts < min_size].index.tolist()
        if not small_clusters:
            break
        # 计算所有簇的中心
        centroids = df.groupby('clustered')[['X', 'Y']].mean()
        # 计算小簇与其他簇的距离
        for small_cluster in small_clusters:
            small_centroid = centroids.loc[small_cluster].values.reshape(1, -1)
            other_clusters = centroids.drop(small_cluster).index
            distances = euclidean_distances(small_centroid, centroids.drop(small_cluster))[0]
            # 找到最近的簇
            closest_cluster = other_clusters[np.argmin(distances)]
            # 合并
            df.loc[df['clustered'] == small_cluster, 'clustered'] = closest_cluster
            break  # 每次合并一个，重新统计
    return df

# 执行调整
df = split_large_clusters(df, MIN_SIZE, MAX_SIZE, TARGET_SIZE)
df = merge_small_clusters(df, MIN_SIZE)

# 重新计算中心
centroids = df.groupby('clustered').agg({'X': 'mean', 'Y': 'mean'}).reset_index()

# 可视化
plt.scatter(df['X'], df['Y'], c=df['clustered'], s=1)
plt.scatter(centroids['X'], centroids['Y'], c='black', s=200)
plt.show()

# 检查最终簇大小
print("最终簇样本数分布：")
print(df['clustered'].value_counts().describe())

二、使用支持样本数约束的替代算法 #

如果不想手动调整，可以直接使用支持簇大小约束的聚类算法，比如sklearn_extra库中的BalancedKMeans：

代码示例

from sklearn_extra.cluster import BalancedKMeans

# 初始化BalancedKMeans，设置每个簇的最小/最大样本数
balanced_kmeans = BalancedKMeans(
    n_clusters=27,
    init='k-means++',
    random_state=0,
    min_cluster_size=MIN_SIZE,
    max_cluster_size=MAX_SIZE
)

df['clustered'] = balanced_kmeans.fit_predict(df[['X', 'Y']])

# 计算中心并可视化
centroids = df.groupby('clustered')[['X', 'Y']].mean().reset_index()
plt.scatter(df['X'], df['Y'], c=df['clustered'], s=1)
plt.scatter(centroids['X'], centroids['Y'], c='black', s=200)
plt.show()

# 检查簇大小
print("Balanced K-Means簇样本数分布：")
print(df['clustered'].value_counts().describe())

关键说明 #

• 层次聚类的核心是基于距离的合并/拆分，天生不支持样本数约束，必须通过后处理调整
• Balanced K-Means是K-Means的改进版，直接支持设置簇的最小/最大样本数，更适配你的需求
• 自定义调整时，拆分大簇用K-Means是因为它能快速将簇划分为大小均衡的子簇，合并小簇基于簇中心距离保证合理性

内容的提问来源于stack exchange，提问作者nxxn