Numpy：使用掩码为不同的行指定不同的值

Question

假设我有一个数组p在形状(m, n)，一个阈值向量Ts在形状(m)。我想使用以下规则替换p中的值：

for i in range(m):
    for j in range(n):
        if p[i, j] > Ts[i]:
            p[i, j] = Ts[i]

我的实施是：

newP = np.zeros_like(p)
cond = p > Ts[:, None]
newP += cond * Ts[:, None]
newP += ~cond * p
p = newP

它看起来确实很难看。我想知道是否有类似p[cond]=Ts风格的方法。谢谢：）

---

例如：

# m = 2, n = 5
p = np.array([[1, 2, 3, 4, 5],
              [6, 7, 8, 9, 10]])
Ts = np.array([3, 8])
expected_new_p = \
    np.array([[1, 2, 3, 3, 3],
              [6, 7, 8, 8, 8]])

Answer 1

您可以通过向Ts添加额外维度来比较p和Ts，提取p < Ts与^{}之间的位置。然后，使用Ts中的值进行OVERWRITE：

i = np.where(p > Ts[:, None])
p[i] = Ts[i[0]]

上面，i是数组的元组，包含p的每个维度中的索引