python - 计算非唯一数组元素的顺序

Question

我正在寻找一种有效的方法来计算每个的“顺序” numpy 数组中的项目，“顺序”定义为项目的数量 前面的元素等于该元素。示例:

order([4, 2, 3, 2, 6, 4, 4, 6, 2, 4])
[0 0 0 1 0 1 2 1 2 3]

当前的解决方案在纯Python中循环并且不够快:

def order(A):
    cnt = defaultdict(int)
    O = np.zeros_like(A)
    for i, r in enumerate(A):
        O[i] = cnt[r]
        cnt[r] += 1
    return O

我使用order来实现scatter:

def scatter(A, c):
    R = A % c
    I = c * order(R) + R
    B = np.full(np.max(I) + 1, -1)
    B[I] = A
    return B

这对于多线程很有用。例如，如果分散 数组包含要写入的地址，然后没有两个线程处理 并行数组将看到相同的地址。

问题是我是否缺少可以使用的 numpy 内置函数 使 order 更快并删除显式循环？

Answer 1

因为你所做的本质上是一个 Pandas cumcount ，并且 Pandas 在内部使用 NumPy，一个想法是看看他们如何实现 cumcount，并做同样的事情。

如果您阅读 Pandas code for cumcount ，内部是这样实现的:

• 对数组进行排序，跟踪每个元素的来源。
• 将已排序数组的每个元素与下一个元素进行比较。如果不同，则是新一轮的开始。 (运行)
• 计算出每组的长度。 (代表)
• 进行累积和，对于不属于新运行的每个元素加 1。 (输出)
• 跟踪每个组受其之前的组影响的程度，该影响不应计算在内。 (out[运行])
• 重复要减去 rep 的值。
• 撤消初始排序，将元素放回原来的位置。

以下是如何在不依赖任何 Pandas 的情况下完成相同的操作。

def order(array):
    # https://github.com/pandas-dev/pandas/blob/v1.3.5/pandas/core/groupby/groupby.py#L1493
    if len(array) == 0:
        return np.array([])
    count = len(array)
    # Can remove 'stable' here to increase speed if you
    # don't care what order the order is assigned in
    ind = np.argsort(array, kind='stable')
    array = array[ind]
    run = np.r_[True, array[:-1] != array[1:]]
    rep = np.diff(np.r_[np.nonzero(run)[0], count])
    out = (~run).cumsum()
    out -= np.repeat(out[run], rep)
    rev = np.empty(count, dtype=np.intp)
    rev[ind] = np.arange(count, dtype=np.intp)
    out = out[rev]
    return out

我发现对于 1000 个元素及更大的数组来说，速度大约快 10 倍。