[英]Element-wise broadcasting for comparing two NumPy arrays?
假設我有一個這樣的數組:
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
。 現在,假設我不想與單個值進行比較,而是想在數組上執行此操作。 換句話說,每個值c
在comparison_array
,找出許多元素如何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
數量級。 那么,我應該如何在numpy
執行此操作以使其快速? 理想情況下,此解決方案應擴展到廣播工作的任何類型的操作,而不僅僅是本示例中的大於或小於。
您可以簡單地向比較數組添加一個維度,以便比較沿着新維度在所有值上“延伸”。
>>> 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])
您將需要轉置用於廣播的數組之一以使其正常工作。 當您一起廣播兩個數組時,維度會對齊,並且任何單位維度都有效地擴展到它們匹配的非單位大小。 所以兩個大小為(16, 1)
(原始數組)和(1, 26)
(比較數組)的數組將廣播到(16, 26)
。
不要忘記對大小為 16 的維度求和:
(base_array[:, None] > comparison_array).sum(axis=1)
切片中的None
等效於np.newaxis
:它是在指定索引處插入新單位維度的眾多方法之一。 你不需要做comparison_array[None, :]
是廣播排列最高維度,並自動填充最低維度。
這是一個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
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.