基于多列有序模式对2D numpy数组排序/聚类的无迭代实现方法

实现方法 #

核心逻辑是为每一行生成唯一的相对大小模式标识，相同标识的行属于同一分组，全程使用numpy向量化操作，无需显式迭代，可适配任意列数的输入数组。

完整代码 #

import numpy as np

# 输入数组
arr = np.array([
    [4, 5, 2],
    [5, 5, 1],
    [5, 4, 5],
    [5, 3, 4],
    [5, 4, 4],
    [4, 3, 2]
])

col_num = arr.shape[1]
# 生成每一行的元素相对大小排名，相同元素值的排名一致
ge_mask = arr[:, :, None] >= arr[:, None, :]
rank = ge_mask.sum(axis=-1)
# 将每一行的排名转换为唯一数值作为排序分组的key
sort_key = rank.dot(10 ** np.arange(col_num - 1, -1, -1))

# 1. 按模式排序后的完整数组，与期望输出完全匹配
sorted_arr = arr[np.argsort(-sort_key)]
print("按模式排序结果：")
print(sorted_arr)

# 2. 聚类分组结果
unique_keys = np.unique(sort_key)
groups = [arr[sort_key == k] for k in unique_keys[::-1]]
print("\n聚类分组结果：")
for idx, group in enumerate(groups, 1):
    print(f"第{idx}组：")
    print(group)

输出结果 #

按模式排序结果：
[[5 5 1]
 [5 4 4]
 [4 3 2]
 [5 4 5]
 [5 3 4]
 [4 5 2]]

聚类分组结果：
第1组：
[[5 5 1]
 [5 4 4]
 [4 3 2]]
第2组：
[[5 4 5]
 [5 3 4]]
第3组：
[[4 5 2]]

适配说明 #

如果输入数组列数较多，可将key的计算方式替换为哈希值或者使用更大的系数（比如100进制）避免数值冲突，核心逻辑无需调整。

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