我需要写一些类似下面的东西。但是如何使用 np.where 或任何其他 vectorization 方法来加速它？
import numpy as np
A = np.arange(5)
B = [4, 3 ,1 ,2]
C = np.zeros(A.shape)
for a, b in zip(A, B):
if C[b] < a - b:
C[b] = a - b
print(C)回答1
警告：如果 B 包含重复元素，这很难（如果可能的话）矢量化循环，因为矢量化 Numpy 操作每次都对整个数组进行操作并且不考虑别名。以下解决方案（如@Ali_Sh 之一）假设没有重复。它还假设 C 与示例中一样为零，并且 A 和 B 的类型相同。
没有重复
比@Ali_Sh 快（两倍）的方法是首先计算差异并使用乘法而不是 np.where 在这种情况下较慢（这是因为 np.where 倾向于在内部使用慢条件分支）。这是代码：
C = np.zeros(A.shape)
tmp = A[:len(B)] - B
C[B] = (tmp > 0) * tmp请注意，使用 B = np.fromiter(B, dtype=A.dtype) 会使代码更快（与其他代码一样），因为 Numpy 每次都将 B 从列表重新转换为数组，而 np.array 并不是最有效的方法。
有重复
如果有重复并且可能发生混叠，那么您可以使用 Numba 来加速操作。与此相反，其他 Numpy 解决方案可能会给出错误的结果。这是代码：
import numba as nb
@nb.njit
def compute_with_aliasing(A, B):
C = np.zeros(A.shape, dtype=A.dtype)
for a, b in zip(A, B):
if C[b] < a - b:
C[b] = a - b
return C
compute_with_aliasing(A, np.fromiter(B, dtype=A.dtype))注意，第一次使用 Numba 函数会因为编译时间的原因而变慢。这种方法在实践中也应该比前一种方法更快。
回答2
你可以这样做：
diff = A[:len(B)] - B # [-4 -2 1 1]
cond = np.less(C[B], diff) # [False False True True]
C[B] = np.where(cond, diff, C[B]) # [0. 1. 1. 0. 0.]@Szczesny 在评论中写了一个解决方案，它很像这个，可以使用 len(B) 而不是 -1 重写为：
C[B] = np.where(A[:len(B)] - B > 0, A[:len(B)] - B, 0)由于与我的答案相似，它可能会被删除。我再次把它放在这里，因为它是单行代码，如果 Szczesny 把它作为一个新的答案，它将被删除。