NumPy 中带有条件的嵌套 for 循环的向量化

Question

我有以下Python代码片段来计算元素

x

和

y

（来自

X

和

Y

）同时验证条件

x<=i

和

y<=j

的次数，其中

i

和

 j

是索引：

import numpy as np

X = np.array([ 8,  7,  9,  5,  3,  2,  4, 10,  6,  1])
Y = np.array([ 1,  3, 10,  9,  5,  7,  8,  2,  6,  4])

n = len(X)

count = 0
expected_count = 0

for i in range(n):
    for j in range(n):
        count += np.sum((X <= (i+1)) & (Y <= (j+1)))
        expected_count += (i + 1) * (j + 1) / n

它工作正常，但出于性能原因，我想向量化此操作。

我知道

i_ = np.arange(1, n+1)
j_ = np.arange(1, n+1)
m_ = i_[:, None] * j_[None, :]

产生一个

n

x

n

矩阵

m_

，其元素总和将返回

expected_count

，并且似乎也是

count

的一种方法。但我缺乏更多的 NumPy 知识来实现这个条件。

我怎样才能实现这一点，特别是在这种情况下矢量化

(X <= (i+1)) & (Y <= (j+1))

操作？

Answer 1

您可以使用：

count = ((X <= i_[:,None]) & (Y  <= j_[:,None])[:,None]).sum()
count
285

expected_count = m_.sum()/n
302.5