首页 > 解决方案 > 使用 1d 列索引数组切片和填充 2d 数组

问题描述

首先,这是我正在尝试做的一维模拟..

假设我有一个 0 的一维数组,我想用 1 替换索引 2 开始的每个 0。我可以这样做:

import numpy as np

x = np.array([0,0,0,0])
i = 2
x[i:] = 1
print(x)  # [0 0 1 1]

现在,我试图找出这个操作的 2d 版本。首先,我有一个 5x4 的 0 数组,例如

foo = np.zeros(shape = (5,4))

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

以及相应的 5 元素列索引数组,例如

fill_locs = np.array([0, 3, 1, 1, 2])

对于 foo 的每一行,我想i:用 1 填充列,其中.i给出的索引是fill_locs. foo[fill_locs.reshape(-1, 1):] = 1感觉不错,但不起作用。

我想要的输出应该看起来像

expected_result = np.array([
    [1, 1, 1, 1],
    [0, 0, 0, 1],
    [0, 1, 1, 1],
    [0, 1, 1, 1],
    [0, 0, 1, 1],
])

标签: pythonnumpy

解决方案


您不需要切片,也不需要创建原始数组。您可以通过广播比较来完成所有这些工作。


a = np.array([0, 3, 1, 1, 2])
n = 4

(a[:, None] <= np.arange(n)).view('i1')

array([[1, 1, 1, 1],
       [0, 0, 0, 1],
       [0, 1, 1, 1],
       [0, 1, 1, 1],
       [0, 0, 1, 1]], dtype=int8)

或使用less_equal.outer

np.less_equal.outer(a, np.arange(n)).view('i1')

推荐阅读