python - 根据单独的 numpy 向量中的值平铺 2D numpy 数组的行

Question

我有一个源数组:

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

还有一个向量，表示我想平铺数组的每一行的次数:

count = array([3, 1, 2])

我想得到:

results =array([[1, 1, 2, 2],
                [1, 1, 2, 2],
                [1, 1, 2, 2],
                [3, 4, 5, 6],
                [7, 7, 7, 8],
                [7, 7, 7, 8]]

是否有矢量化/numpy 方法来实现这一点？

目前我正在使用迭代循环方法，当 len(a) 和/或 count 包含高值时，它的速度非常慢。

Answer 1

numpy.repeat()是你所追求的:

代码:

np.repeat(a, count, axis=0)

测试代码:

import numpy as np

a = np.array([[1, 1, 2, 2],
              [3, 4, 5, 6],
              [7, 7, 7, 8]])

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

print(np.repeat(a, count, axis=0))

结果:

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