根据一维计数器数组填充二维数组列

Posted 2023-02-24

【中文标题】根据一维计数器数组填充二维数组列

【英文标题】：Filling of 2D array columns according to 1D counter array

【发布时间】：2018-05-26 22:28:50

【问题描述】：

我正在寻找一种 numpy 解决方案，用不同的一维计数器数组（在下面的例子）。

我尝试了以下方法：

import numpy as np

cnt = np.array([1, 3, 2, 4])   # cnt array: how much elements per column are 1
a = np.zeros((5, 4))           # array that must be filled with 1s per column
for i in range(4):             # for each column
    a[:cnt[i], i] = 1          # all elements from top to cnt value are filled

print(a)

并给出所需的输出：

[[1. 1. 1. 1.] 
 [0. 1. 1. 1.]
 [0. 1. 0. 1.]
 [0. 0. 0. 1.]
 [0. 0. 0. 0.]]

有没有一种更简单（和更快）的方法可以使用 numpy 例程来执行此操作，而无需每列循环？

a = np.full((5, 4), 1, cnt)

类似上面的东西会很好，但不起作用。

感谢您的宝贵时间！

【参考方案1】：

您可以使用np.where 并像这样进行广播：

>>> import numpy as np
>>> 
>>> cnt = np.array([1, 3, 2, 4])   # cnt array: how much elements per column are 1
>>> a = np.zeros((5, 4))           # array that must be filled with 1s per column
>>> 
>>> res = np.where(np.arange(a.shape[0])[:, None] < cnt, 1, a)
>>> res
array([[1., 1., 1., 1.],
       [0., 1., 1., 1.],
       [0., 1., 0., 1.],
       [0., 0., 0., 1.],
       [0., 0., 0., 0.]])

或就地：

>>> a[np.arange(a.shape[0])[:, None] < cnt] = 1
>>> a
array([[1., 1., 1., 1.],
       [0., 1., 1., 1.],
       [0., 1., 0., 1.],
       [0., 0., 0., 1.],
       [0., 0., 0., 0.]])

【讨论】：

感谢您的精彩回答：比我需要的更强大。 “res”二维数组仅根据“cnt”条件将单元格更改为特定数字（在本例中为 1），并在所有其他单元格中保存“a”的原始值。这允许多次使用一维“cnt”数组动态调整二维“res”数组。作为一名 Python 初学者，我在比较一维垂直数组“np.arange(a.shape[0])[:, None]”（本例中为 5 个元素）与一维水平数组“cnt”（4在这种情况下的元素）用于重建所需的二维“res”数组。谢谢！

@wilmert 该条件利用广播。比较形状 5x1 和 4 的数组会得到以下结果：1) 通过在 左侧插入轴 -> 1x4 将操作数 2 提升为 ndim 2 2) 通过沿重复方向将 5x1 和 1x4 都提升为 5x4放大的轴（仅在概念上，实际上没有数据被复制） 3）实际的元素比较完成了一个布尔 5x4 数组。