给定行索引和起始列索引，如何将值分配给 numpy 数组？答案

【问题标题】：how to assign values to numpy array given the row indices and starting column indices?给定行索引和起始列索引，如何将值分配给 numpy 数组？
【发布时间】：2021-02-07 02:30:25
【问题描述】：

a = np.array([2,3,1,4])
b = np.array([2,3,7,1])
c = np.zeros((4, 10))

我想将值 1 分配给 c 中的某些元素。 a 和 b 定义了这些元素的位置。 a 是每行中值 1 的起始列索引。而b 表示一行中有多少个连续的1。我期待的输出是：

array([[ 0.,  0.,  1.,  1.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  1.,  1.,  1.,  0.,  0.,  0.,  0.],
       [ 0.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.]])

我可以使用如下简单的 for 循环：

for i in range(c.shape[0]):
    for k in range(a[i], a[i]+b[i]):
        c[i,k]=1

但是对于大型数组来说会很慢，有没有更快的 numpy 索引来做到这一点？谢谢。

【问题讨论】：

c[i, a[i]:a[i]+b[i]] = 1 应该可以代替 k 循环工作。
是的，我主要关心的是最外层的 for 循环。
如果不先处理内部循环，就无法处理外部循环。
c 在您的实际用例中的典型形状是什么？

标签： python numpy numpy-ndarray numpy-slicing

【解决方案1】：

我实现了下一个解决方案，没有任何 Python 循环，只是纯 NumPy 代码。也许它不像 python-loop 解决方案那么简单，但肯定会更快，尤其是对于大数据。

Try it online!

import numpy as np

def set_val_2d(a, val, starts, lens):
    begs = starts + np.arange(a.shape[0]) * a.shape[1]
    ends = begs + lens
    clens = lens.cumsum()
    ix = np.ones((clens[-1],), dtype = np.int64)
    ix[0] = begs[0]
    ix[clens[:-1]] = begs[1:] - ends[:-1] + 1
    ix = ix.cumsum()
    a.ravel()[ix] = val

a = np.array([2,3,1,4])
b = np.array([2,3,7,1])
c = np.zeros((4, 10))

set_val_2d(c, 1, a, b)
print(c)

输出：

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

【讨论】：

【解决方案2】：

您可以将其转换为一维问题

def convert_inds(a,b,array_shape):
    
    nrows,ncols = array_shape
    to_take = np.zeros(sum(b))
    count = 0
    for ind,item in enumerate(b):
        start_ind = ind*ncols+a[ind]
        to_take[count:count+item] = np.arange(start_ind,start_ind+item)
        count += item
        
    return to_take.astype(np.int)

to_take = convert_inds(a,b,c.shape)

c.ravel()[to_take] = 1

在上面的代码中，convert_inds 会将a 和b 转换为

array([ 2,  3, 13, 14, 15, 21, 22, 23, 24, 25, 26, 27, 34])

它们是扁平化c 中1s 的索引。通过这样做，您只需要在函数convert_inds 中迭代b。

【讨论】：

【解决方案3】：

如果你选择一种基于索引的方法，最困难的部分是找到轴 1 的索引。这非常类似于：

>>> np.repeat(a, b)
array([2, 2, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 4])

除了每组索引应该递增。可以使用此功能完成此修复：

def accumulative_count(counts, initial):
    counter = np.ones(np.sum(counts), dtype=int)
    marker_idx = np.r_[0, np.cumsum(counts)[:-1]]
    subtract_vals = np.r_[1, counts[:-1]]
    initial_vals = np.r_[initial[0], np.diff(initial)]
    counter[marker_idx] = counter[marker_idx] - subtract_vals + initial_vals
    return np.cumsum(counter)

>>> accumulative_count(counts, initial)
array([2, 3, 3, 4, 5, 1, 2, 3, 4, 5, 6, 7, 4], dtype=int32)

毕竟，你有能力完成它：

c[np.repeat(np.arange(len(c)), b), accumulative_count(b, a)] = 1

c:

array([[0., 0., 1., 1., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 1., 1., 1., 0., 0., 0., 0.],
       [0., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
       [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.]])

【讨论】：