python - 从数组值动态填充 numpy 矩阵？

Question

我正在尝试根据数组的值使用 numpy 动态构造一个二维矩阵，如下所示:

In [113]: A = np.zeros((5,5),dtype=bool)
In [114]: A
Out[114]: array([[False, False, False, False, False],
   [False, False, False, False, False],
   [False, False, False, False, False],
   [False, False, False, False, False],
   [False, False, False, False, False]], dtype=bool)

 In [116]: B = np.array([0,1,3,0,2])

 In [117]: B
 Out[117]: array([0, 1, 3, 0, 2])

现在，我想使用 B 的值将每行的前 n 个值分配给 A 为 True。对于这个 A 和 B，正确的输出是:

In [118]: A
Out[118]: array([[False, False, False, False, False],
   [ True, False, False, False, False],
   [ True,  True,  True, False, False],
   [False, False, False, False, False],
   [ True,  True, False, False, False]], dtype=bool)

B的长度将始终等于A的行数，并且B的值将始终小于或等于A的列数。A的大小和B的值不断变化，所以我需要即时构建这些。

我确信这在 numpy 中有一个简单的解决方案，但我花了最后一个小时来研究重复、平铺和我能想到的任何其他东西的变化。在我遭受脑震荡之前，有人可以帮助我吗？ :)

编辑:我需要经常这样做，所以速度将是一个问题。我现在能想到的唯一版本是这样的:

np.vstack([ [True]*x + [False]*(500-x) for x in B ])

但我预计由于 for 循环，这会很慢(如果我有什么可以比较的话，我会计时)。

Answer 1

怎么样:

>>> A = np.zeros((5, 7),dtype=bool)
>>> B = np.array([0,1,3,0,2])
>>> (np.arange(len(A[0])) < B[:,None])
array([[False, False, False, False, False, False, False],
       [ True, False, False, False, False, False, False],
       [ True,  True,  True, False, False, False, False],
       [False, False, False, False, False, False, False],
       [ True,  True, False, False, False, False, False]], dtype=bool)

(我从 (5,5) 更改了形状，因为我对哪个轴是哪个轴感到困惑，并且我想确保我使用的是正确的轴。)

[简化自 (np.arange(len(A[0]))[:,None] < B).T -- 如果我们展开B而不是A ，不需要转置。]