如何在Python中统计二维数组每行的唯一值数量？

问题示例 #

输入数组：

import numpy as np

arr = np.array([[3,4,4,4,3,4], 
                [4,4,4,4,4,4], 
                [3,3,3,2,3,2],
                [2,3,3,1,2,2]])

期望输出：

res = np.array([2,1,2,3])

解决方案 #

方法一：简洁写法（适用于小型数组）

利用np.apply_along_axis对每行调用np.unique统计唯一值数量：

res = np.apply_along_axis(lambda x: len(np.unique(x)), axis=1, arr=arr)

该方法逻辑直观，一行代码即可实现，但本质是逐行循环处理，大数据量下效率一般。

方法二：向量化高效实现（适用于大型数组）

通过排序+差分的方式避免循环，利用NumPy底层优化提升处理效率：

# 对每行排序，相同元素会相邻排列
sorted_arr = np.sort(arr, axis=1)
# 计算每行相邻元素的差值
diff = np.diff(sorted_arr, axis=1)
# 统计每行非零差值的数量，加1即为唯一值个数（n个不同元素对应n-1个非零差值）
res = np.sum(diff != 0, axis=1) + 1

方法三：基于np.bincount的实现（仅适用于非负整数元素）

如果数组元素均为非负整数，可以用np.bincount统计每个值的出现次数，再统计非零次数：

res = np.array([np.count_nonzero(np.bincount(row)) for row in arr])

注意：若行中存在负数，np.bincount会抛出错误，此时不适用该方法。

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