如何使用填充沿新维度重复numpy数组？

Question

给定一个两个数组：一个输入数组和一个重复数组，我想接收一个数组，它沿着一个新维度对每一行重复指定的次数，并填充到最后。你知道吗

to_repeat = np.array([1, 2, 3, 4, 5, 6])
repeats = np.array([1, 2, 2, 3, 3, 1])
# I want final array to look like the following:
#[[1, 0, 0],
# [2, 2, 0],
# [3, 3, 0],
# [4, 4, 4],
# [5, 5, 5],
# [6, 0, 0]]

问题是，我使用的是大型数据集（大约10万个），所以列表理解太慢了-实现这一点的快速方法是什么？你知道吗

Answer 1

这是一个基于^{}的masking

m = repeats[:,None] > np.arange(repeats.max())
out = np.zeros(m.shape,dtype=to_repeat.dtype)
out[m] = np.repeat(to_repeat,repeats)

样本输出-

In [44]: out
Out[44]: 
array([[1, 0, 0],
       [2, 2, 0],
       [3, 3, 0],
       [4, 4, 4],
       [5, 5, 5],
       [6, 0, 0]])

或广播乘法-

In [67]: m*to_repeat[:,None]
Out[67]: 
array([[1, 0, 0],
       [2, 2, 0],
       [3, 3, 0],
       [4, 4, 4],
       [5, 5, 5],
       [6, 0, 0]])

对于大型数据集/大小，我们可以利用multi-cores并在broadcasting上使用^{} module提高内存效率-

In [64]: import numexpr as ne

# Re-using mask `m` from previous method
In [65]: ne.evaluate('m*R',{'m':m,'R':to_repeat[:,None]})
Out[65]: 
array([[1, 0, 0],
       [2, 2, 0],
       [3, 3, 0],
       [4, 4, 4],
       [5, 5, 5],
       [6, 0, 0]])