在 NumPy 数组上计算 "moving sum of counts"

Computing a "moving sum of counts" on a NumPy array

我有以下数组:

# input
In [77]: arr = np.array([23, 45, 23, 0, 12, 45, 45])

# result
In [78]: res = np.zeros_like(arr)

现在,我想计算唯一元素的 移动总和 并将其存储在 res 数组中。

具体来说,res数组应该是:

In [79]: res
Out[79]: array([1, 1, 2, 1, 1, 2, 3])

[23, 45, 23, 0, 12, 45, 45]
[1,   1,   2,   1,   1,   2,   3]

我们开始对每个元素进行计数,如果元素重新出现,则增加 count,直到到达数组末尾。此元素特定计数应作为结果返回。


我们应该如何使用 NumPy 内置函数来实现这一点?我尝试使用 numpy.bincount 但它给出了不希望的结果。

不确定您是否会找到内置函数,所以这是一个使用 argsort 的自制程序。

def running_count(arr):
    idx = arr.argsort(kind='mergesort')
    sarr = arr[idx]
    neq = np.where(sarr[1:] != sarr[:-1])[0] + 1
    run = np.ones(arr.shape, int)
    run[neq[0]] -= neq[0]
    run[neq[1:]] -= np.diff(neq)
    res = np.empty_like(run)
    res[idx] = run.cumsum()
    return res

例如:

>>> running_count(arr)
array([1, 1, 2, 1, 1, 2, 3])
>>> running_count(np.array(list("xabaaybeeetz")))
array([1, 1, 1, 2, 3, 1, 2, 1, 2, 3, 1, 1])

解释器:

我们首先使用 argsort 进行排序,因为我们需要索引最终返回到原始顺序。这里有一个稳定的排序很重要,因此使用慢速合并排序。

元素排序后,运行 计数将形成 "saw tooth" 模式。创建它的矢量化方法是观察锯齿的差异具有 "jump" 值,新齿开始的位置和其他任何地方的值。所以这很容易构建。