Python Numpy - 创建长度基于一维数组的二维数组答案

【问题标题】：Python Numpy - Create 2d array where length is based on 1D arrayPython Numpy - 创建长度基于一维数组的二维数组
【发布时间】：2022-01-07 04:16:53
【问题描述】：

抱歉标题混乱，但不知道如何使它更简洁。这是我的要求：

arr1 = np.array([3,5,9,1])
arr2 = ?(arr1)

arr2 将是：

[
[0,1,2,0,0,0,0,0,0],
[0,1,2,3,4,0,0,0,0],
[0,1,2,3,4,5,6,7,8],
[0,0,0,0,0,0,0,0,0]
]

它不需要根据最大值而变化，形状是预先知道的。所以开始我已经能够得到一个零的形状：

arr2 = np.zeros((len(arr1),max_len))

然后我当然可以像这样在 arr1 上做一个 for 循环：

for i, element in enumerate(arr1):
    arr2[i,0:element] = np.arange(element)

但这可能需要很长时间，而且这里的两个维度都相当大（arr1 是几百万行，max_len 大约是 500）。在 numpy 中是否有一种干净优化的方法来做到这一点？

【问题讨论】：

你能通过对np.arange(10) 和arr1 进行某种“外部”操作来创建一个“掩码”吗？ True 序列为 arr1?

标签： python numpy vectorization

【解决方案1】：

基于几年前@Divakar 发布的“填充”想法：

In [161]: res = np.arange(9)[None,:].repeat(4,0)
In [162]: res[res>=arr1[:,None]] = 0
In [163]: res
Out[163]: 
array([[0, 1, 2, 0, 0, 0, 0, 0, 0],
       [0, 1, 2, 3, 4, 0, 0, 0, 0],
       [0, 1, 2, 3, 4, 5, 6, 7, 8],
       [0, 0, 0, 0, 0, 0, 0, 0, 0]])

【讨论】：

【解决方案2】：

用itertools.zip_longest试试这个-

import numpy as np
import itertools

l = map(range, arr1)
arr2 = np.column_stack((itertools.zip_longest(*l, fillvalue=0)))
print(arr2)

array([[0, 1, 2, 0, 0, 0, 0, 0, 0],
       [0, 1, 2, 3, 4, 0, 0, 0, 0],
       [0, 1, 2, 3, 4, 5, 6, 7, 8],
       [0, 0, 0, 0, 0, 0, 0, 0, 0]])

【讨论】：

zip_longest 是“填充”lsts 或数组的好工具。虽然有一种使用数组“掩码”填充的巧妙方法，但这种迭代工具更容易记住和应用。

【解决方案3】：

我在@hpaulj 的答案上添加了一些细微的变化，因为您提到max_len 在500 附近并且您有数百万行。在这种情况下，您可以预先计算一个 500 x 500 的矩阵，其中包含所有可能的行并使用 arr1 对其进行索引：

import numpy as np
np.random.seed(0)

max_len = 500
arr = np.random.randint(0, max_len, size=10**5)

# generate all unique rows first, then index
# can be faster if max_len << len(arr)
# 53 ms
template = np.tril(np.arange(max_len)[None,:].repeat(max_len,0), k=-1)
res = template[arr,:]

# 173 ms
res1 = np.arange(max_len)[None,:].repeat(arr.size,0)
res1[res1>=arr[:,None]] = 0

assert (res == res1).all()

【讨论】：