NumPy数组向量化累积计数的高效实现方案问询

嘿，你的这个问题我之前也碰到过，先给你拆解清楚～

首先回应你关于稀疏矩阵的疑问：是的，用稀疏矩阵做行累积和后肯定会变成稠密矩阵。原因很简单，你原本在每行只有一个位置是1，累积和之后，从那个位置开始的所有列都会变成1（因为累积是逐步加总）。如果你的数组里数值范围很大（比如最大数是10000），那这个矩阵的内存占用会直接爆炸，完全得不偿失，所以这个思路确实不太可行。

接下来给你几个更高效的实现方式，都是纯NumPy或者常用工具库的方法：

方法1：基于np.unique和分组累积（简单直观） #

这个方法适合唯一值不多的场景，逻辑清晰易懂：

import numpy as np

# 示例数组
arr = np.array([2, 2, 3, 2, 3])
# 获取唯一值、逆索引（每个元素对应唯一值的下标）
_, idx = np.unique(arr, return_inverse=True)

result = np.zeros_like(arr)
# 对每个唯一值对应的位置做累积计数
for val_idx in np.unique(idx):
    mask = idx == val_idx
    result[mask] = np.cumsum(np.ones_like(result[mask]))

print(result)  # 输出: [1 2 1 3 2]

还能更简洁一点，直接对布尔掩码做累积：

result = np.zeros_like(arr)
for val in np.unique(arr):
    mask = arr == val
    result[mask] = np.cumsum(mask)

方法2：基于排序的无循环高效实现（适合大数据量） #

如果你的数组很大，循环会拖慢速度，这个纯NumPy的无循环方法效率拉满：

import numpy as np

arr = np.array([2, 2, 3, 2, 3])
# 先按数组值排序，得到排序后的索引
sorted_indices = np.argsort(arr)
sorted_arr = arr[sorted_indices]

# 初始化累积计数
cum_counts = np.cumsum(np.ones_like(sorted_arr))
# 找到每个新值的起始位置，重置计数为1
value_changes = np.diff(sorted_arr, prepend=-np.inf) != 0
cum_counts[value_changes] = 1

# 将计数映射回原数组的位置
result = np.zeros_like(arr)
result[sorted_indices] = cum_counts

print(result)  # 输出: [1 2 1 3 2]

这个方法利用排序把相同值聚在一起，一次性计算累积计数，再映射回原位置，完全避免了循环，大数据量下性能提升特别明显。

方法3：用Pandas快速实现（如果允许用Pandas） #

如果你项目里已经在用Pandas，那一行代码就能搞定，省心又方便：

import pandas as pd
import numpy as np

arr = np.array([2, 2, 3, 2, 3])
result = pd.Series(arr).groupby(arr).cumcount() + 1  # +1是因为cumcount从0开始
print(result.to_numpy())  # 输出: [1 2 1 3 2]

总结一下：稀疏矩阵的思路确实不推荐，会导致严重的内存浪费；优先选择纯NumPy的无循环方法（方法2），如果追求代码简洁可以用方法1或Pandas的方法。

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