2维NumPy数组中每个元素的计数

Question

假设您有一个二维数组（作为NumPy的int数组），例如：

[[2,2,3,3],
 [2,3,3,3],
 [3,3,4,4]]

现在你想要一个与原始数组形状相同的数组，但是不是使用原始值，而是用其出现次数来替换数字。这意味着数字2变为3，因为它出现了3次，3变为7，4变为2。

所以输出将是：

[[3,3,7,7],
 [3,7,7,7],
 [7,7,2,2]]

我的解决方案是首先创建一个字典，将所有原始值保存为键，并将出现次数保存为值。但对于形状为2000x2000的数组，这似乎非常慢。

我如何更有效地实现这个呢？

谢谢！

- Tschoko Kuki

---

在你的例子中，数组的扁平化版本是已排序的。这种情况是否总是成立？ - Brad Solomon

不，这个数组不是用来排序的。谢谢！ - Tschoko Kuki

Answer 1

我相信你可以通过在np.unique()中使用return_inverse来保留NumPy:

如果为True，则还返回唯一数组的索引（对于提供的轴），这些索引可用于重构ar。

>>> import numpy as np

>>> a = np.array([[2,2,3,3],
...               [2,3,3,3],
...               [3,3,4,4]])

>>> _, inv, cts = np.unique(a, return_inverse=True, return_counts=True)
>>> cts[inv].reshape(a.shape)

array([[3, 3, 7, 7],
       [3, 7, 7, 7],
       [7, 7, 2, 2]])

这也适用于扁平化数组未排序的情况，例如b = np.array([[1, 2, 4], [4, 4, 1]])。

- Brad Solomon

Answer 2

一种方法是使用 numpy.unique 提取值的计数。

然后将其转换为字典，并使用 numpy.vectorize 来利用这个字典映射。

import numpy as np

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

d = dict(zip(*np.unique(A.ravel(), return_counts=True)))

res = np.vectorize(d.get)(A)

array([[3, 3, 7, 7],
       [3, 7, 7, 7],
       [7, 7, 2, 2]], dtype=int64)

性能

我发现上述方法对于一个2000x2000的数组需要大约2秒，而使用基于collections.Counter字典的方法需要3秒。但是PaulPanzer和BradSolomon提出的纯numpy解决方案仍然更快。

import numpy as np
from collections import Counter

A = np.random.randint(0, 10, (2000, 2000))
MAX_LOOKUP = 2**24

def map_count(A):
    d = dict(zip(*np.unique(A.ravel(), return_counts=True)))
    return np.vectorize(d.get)(A)

def map_count2(A):
    d = Counter(A.ravel())
    return np.vectorize(d.get)(A)

def bs(A):
    _, inv, cts = np.unique(A, return_inverse=True, return_counts=True)
    return cts[inv].reshape(A.shape)

def pp(a):
    mn, mx = a.min(), a.max()
    span = mx-mn+1
    if span > MAX_LOOKUP:
        raise RuntimeError('values spread to wide')
    a = a - mn
    return np.bincount(a.ravel(), None, span)[a]

%timeit map_count(A)   # 1.9 s ± 24.2 ms per loop
%timeit map_count2(A)  # 3 s ± 33.1 ms per loop
%timeit bs(A)          # 887 ms ± 20 ms per loop
%timeit pp(A)          # 149 ms ± 6.32 ms per loop

- jpp

---

我以前从未听说过 np.vectorize。它听起来相当不错。 - Mateen Ulhaq

听起来不错，但并非如此。因为它实际上并没有"向量化"——它只是循环遍历每个项目并应用一个函数。更多的是方便而非真正提高性能。 - jpp

由于某种原因，我认为它会查看代码图并正确地进行矢量化。 - Mateen Ulhaq

Answer 3

这里有一种方法，利用了您的值是 int 类型的事实：

MAX_LOOKUP = 2**24

def f_pp(a):
    mn, mx = a.min(), a.max()
    span = mx-mn+1
    if span > MAX_LOOKUP:
        raise RuntimeError('values spread to wide')
    a = a - mn
    return np.bincount(a.ravel(), None, span)[a]

时间（大部分基于@jpp的工作）：

>>> from timeit import timeit
>>> kwds = dict(globals=globals(), number=3)
>>> 
>>> for l, r in [(0, 10), (0, 1000), (-8000000, 8000000)]:
...     a = np.random.randint(l, r, (2000, 2000))
...     print(l, r)
...     print('mc ', timeit('map_count(a)', **kwds))
...     print('mc2', timeit('map_count2(a)', **kwds))
...     print('bs ', timeit('bs(a)', **kwds))
...     print('pp ', timeit('f_pp(a)', **kwds))
... 
0 10
mc  2.462232475867495
mc2 3.820418732939288
bs  1.266723491018638
pp  0.11216754489578307
0 1000
mc  2.972961534978822
mc2 4.3769155589398
bs  2.1607728030066937
pp  0.14146877988241613
-8000000 8000000
mc  10.753600731957704
mc2 8.373655589064583
bs  2.700256273150444
pp  0.7070535880047828

- Paul Panzer

---

你能简要解释一下为什么我们会看到这个算法的巨大改进吗？ - jpp

@jpp 只要箱子的数量不是过于庞大，bincount 的时间复杂度为 O(n)，在我的经验中它是一个非常快速的函数。unique 在内部进行排序，因此时间复杂度为 O(n log n)。 - Paul Panzer

谢谢 - 这个答案值得被采纳！ - jpp