基于另一个 numpy 向量的 numpy 矩阵条目的向量化/高效部分答案

【问题标题】：Vectorized/efficient section of numpy matrix-entries based on another numpy vector基于另一个 numpy 向量的 numpy 矩阵条目的向量化/高效部分
【发布时间】：2026-02-07 17:55:01
【问题描述】：

以下是我想要找到有效解决方案的情况示例：

import numpy as np
a = np.array([[1,2,3],[4,5,6],[7,8,9]])
b = np.array([0,1,2])

b 向量的每一行应该有一个整数条目，范围在 0 到 a 的列数之间。所需的输出 c 是一个形状为 a 的矩阵，第 i 行包含每个条目直到（并且不包括）a 的第 i 行的第 b[i] 个条目，该行的其余部分应该用零填充。因此，在给出的示例中，我们寻找以下解决方案：

c = np.array([[0,0,0],[4,0,0],[7,8,0]])

有两种“简单”的方法可以做到这一点：

c = np.zeros(a.shape)
for i in range(a.shape[0]):
    c[i, :b[i]] =  a[i, :b[i]]

另一种方法是首先定义一个辅助矩阵，然后针对a进行广播。

aux = np.zeros(a.shape)
for i in range(a.shape[0]):
    aux[i, :b[i]] = 1 
c = a * aux

我正在寻找的是 矢量化解决方案，当我增加 a 的行数时，它们可以在运行时很好地扩展。

【问题讨论】：

这有帮助吗：*.com/a/67173501/7831421。请注意，在此答案中 use = b[:, None] 和 have 是您的 a
感谢您的提示。使用该链接，我想出了以下想法：首先是n, m = a.shape，然后是aux = np.repeat(np.arange(m)[None,:], n, axis = 0) < b，最后是c = a * aux。但是，我在 python 方面的知识太有限，无法判断当我增加行数时，这是否是运行时扩展性最好的解决方案。
是的。如果您在创建aux 时使用b[:, None] 代替b，那应该可以完成这项工作。
@H1ghfiv3 您需要重塑 b 以便正确广播：aux = np.repeat(np.arange(m)[None,:], n, axis = 0) < b.reshape(-1, 1)（尝试使用非方阵 a）
对，我将 b 的输入形状与我在我的电脑上测试过代码的（正确的）形状混合在一起。谢谢

标签： python python-3.x numpy

【解决方案1】：

作为一个选项，您可以制作一个掩码矩阵，并将 a 元素乘以该掩码：

a * (np.mgrid[:a.shape[0], :a.shape[1]][1] < b.reshape(-1, 1))

输出：

array([[0, 0, 0],
       [4, 0, 0],
       [7, 8, 0]])

下面是如何构建掩码的分步说明：

np.mgrid[:a.shape[0], :a.shape[1]][1]
# array([[0, 1, 2],
#        [0, 1, 2],
#        [0, 1, 2]])

np.mgrid[:a.shape[0], :a.shape[1]][1] < b.reshape(-1, 1)
# array([[False, False, False],
#        [ True, False, False],
#        [ True,  True, False]])

【讨论】：

【解决方案2】：

使用 numpy 广播为 indices > b[i] 创建一个布尔值 mask，即 False，即，a 应替换为 0 的索引。然后将a 与mask 相乘得到c。

n, m = a.shape
mask = np.repeat(np.arange(m)[None, :], n, axis = 0) < b[:, None]
c = mask * a

输出：

>>> c
array([[0, 0, 0],
       [4, 0, 0],
       [7, 8, 0]])

【讨论】：