Question

我想从numpy数组的行中获取前n个值，其中n是在单独的1-d数组中指定的：

import numpy as np

a = np.zeros((5, 5))
n = [1, 3, 2, 4, 1]

result = [[1, 0, 0, 0, 0],
          [1, 1, 1, 0, 0],
          [1, 1, 0, 0, 0],
          [1, 1, 1, 1, 0],
          [1, 0, 0, 0, 0]]

我正在寻找一种不需要迭代的解决方案，因为结果数组将具有数百万行。

Answer 1

在利用broadcasting创建遮罩以及最终数组时，对范围数组使用n的外部比较-

ncols = 5
mask_out = np.greater.outer(n,np.arange(ncols))

样品运行-

In [19]: n = [1, 3, 2, 4, 1]

In [9]: ncols = 5

# Output as mask
In [10]: np.greater.outer(n,np.arange(ncols))
Out[10]: 
array([[ True, False, False, False, False],
       [ True,  True,  True, False, False],
       [ True,  True, False, False, False],
       [ True,  True,  True,  True, False],
       [ True, False, False, False, False]])

# Output as array of 0s and 1s
In [11]: np.greater.outer(n,np.arange(ncols)).view('i1')
Out[11]: 
array([[1, 0, 0, 0, 0],
       [1, 1, 1, 0, 0],
       [1, 1, 0, 0, 0],
       [1, 1, 1, 1, 0],
       [1, 0, 0, 0, 0]], dtype=int8)

如果您必须填充已经初始化的数组result，只需使用mask_out来掩盖该数组，即result[mask_out] = ...。

Answer 2

这是一种使用高级索引到模板行数组中的方法：

def f_pp(a, ncol=None):
    if ncol is None:
        ncol = a.max()
    t = np.array([1, 0], 'u1').repeat(ncol)
    ts, = t.strides
    t = np.lib.stride_tricks.as_strided(t[ncol:], (ncol+1, ncol), (-ts, ts))
    return t[a]

使用@Divakar作为参考的时间：

def f_div(a, ncol=None):
    if ncol is None:
        ncol = a.max()
    return np.greater.outer(a, np.arange(ncol))

from timeit import timeit

for ncol, nrow in [(10, 10**7), (100, 10**6), (1000, 10**5), (10000, 10000)]:
    a = np.random.randint(0, ncol+1, nrow)
    print('\n', ncol, nrow, ':')
    print('div', timeit(lambda: f_div(a), number=10))
    print('pp ', timeit(lambda: f_pp(a), number=10))

打印：

 10 10000000 :
div 2.4297873955219984
pp  1.698299034498632

 100 1000000 :
div 1.465646200813353
pp  0.4803247870877385

 1000 100000 :
div 1.3471891237422824
pp  0.35979613568633795

 10000 10000 :
div 0.6783521044999361
pp  0.38309483136981726

Answer 3

我认为您正在尝试获得类似的东西？我可能不知道有一种numpy特定的方法会更快，但是如果速度不是问题，就足够了：

import numpy as np

a = np.zeros((5, 5))
n = [1, 3, 2, 4, 1]

for index in range(len(n)):
    a[index][0:n[index]] = 1

print(a)

输出

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

从一维NumPy数组创建二维蒙版

3 个答案: