针对两个NumPy数组进行逐元素广播比较的方法？

Question

假设我有一个这样的数组：

import numpy as np

base_array = np.array([-13, -9, -11, -3, -3, -4,   2,  2,
                         2,  5,   7,  7,  8,  7,  12, 11])

假设我想知道：base_array中有多少个元素大于4？这可以通过利用广播来简单实现：

np.sum(4 < base_array)

答案是7。现在，假设我想要对一个数组进行比较，而不是与单个值进行比较。换句话说，对于comparison_array中的每个值c，找出有多少个base_array元素大于c。如果我以天真的方式这样做，显然会失败，因为它不知道如何正确广播：

comparison_array = np.arange(-13, 13)
comparison_result = np.sum(comparison_array < base_array)

输出：

Traceback (most recent call last):
  File "<pyshell#87>", line 1, in <module>
    np.sum(comparison_array < base_array)
ValueError: operands could not be broadcast together with shapes (26,) (16,) 

如果我可以让comparison_array的每个元素都广播到base_array的形状上，那么就能解决这个问题。但我不知道如何进行“逐元素广播”。

现在，我知道如何使用列表推导为两种情况都实现这一点:

first = sum([4 < i for i in base_array])
second = [sum([c < i for i in base_array])
          for c in comparison_array]
print(first)
print(second)

输出：

7
[15, 15, 14, 14, 13, 13, 13, 13, 13, 12, 10, 10, 10, 10, 10, 7, 7, 7, 6, 6, 3, 2, 2, 2, 1, 0]

但是众所周知，在处理大型数组时，这种方法比正确进行向量化的numpy实现慢了数个数量级。那么，我应该如何在numpy中快速完成此操作呢？理想情况下，这种解决方案应该适用于任何广播可行的操作，而不仅仅是在此示例中的大于或小于运算。

- dain

Answer 1

你可以简单地向比较数组添加一个维度，这样比较就会在新维度上“拉伸”到所有值。

>>> np.sum(comparison_array[:, None] < base_array)
228

这是广播的基本原理，适用于所有类型的操作。

如果您需要沿着轴计算总和，只需在比较后指定要求和的轴即可。

>>> np.sum(comparison_array[:, None] < base_array, axis=1)
array([15, 15, 14, 14, 13, 13, 13, 13, 13, 12, 10, 10, 10, 10, 10,  7,  7,
        7,  6,  6,  3,  2,  2,  2,  1,  0])

- miradulo

---

我需要答案的形状与“comparison_array”相同，就像列表推导式所给出的那样。 - dain

Answer 2

你需要转置其中一个数组才能正确地进行广播。当你将两个数组一起广播时，它们的维度会对齐，并且任何单位维度都会被有效地扩展到它们匹配的非单位大小。因此，大小为（16，1）（原始数组）和（1，26）（比较数组）的两个数组将广播到（16，26）。

不要忘记沿着大小为16的维度进行求和：

(base_array[:, None] > comparison_array).sum(axis=1)

None在切片中等同于np.newaxis：这是在指定索引处插入新的单元维度的众多方法之一。你不需要执行comparison_array[None, :]的原因是广播会对齐最高维度，并自动用1填充最低维度。

- Mad Physicist

Answer 3

这里有一个关注内存效率和性能的例子，涉及到 np.searchsorted -

def get_comparative_sum(base_array, comparison_array):
    n = len(base_array)
    base_array_sorted = np.sort(base_array)
    idx = np.searchsorted(base_array_sorted, comparison_array, 'right')
    idx[idx==n] = n-1
    return n - idx - (base_array_sorted[idx] == comparison_array)

计时 -

In [40]: np.random.seed(0)
    ...: base_array = np.random.randint(-1000,1000,(10000))
    ...: comparison_array = np.random.randint(-1000,1000,(20000))

# @miradulo's soln
In [41]: %timeit np.sum(comparison_array[:, None] < base_array, axis=1)
1 loop, best of 3: 386 ms per loop

In [42]: %timeit get_comparative_sum(base_array, comparison_array)
100 loops, best of 3: 2.36 ms per loop

- Divakar