重复值 n 次，n 在一个数组中

Question

我有一个 a 值数组，其值在每个索引处 idx 我想重复一定数量 b[idx] 次，在同一索引处给定 idx 在另一个数组 (b) 中，像这样：

a = numpy.array([1, 2, 3 ,4, 5])
b = numpy.array([2, 3, 1, 2, 4])

期望的输出：

c = numpy.array([1, 1, 2, 2, 2, 3, 4, 4, 5, 5, 5, 5])

我意识到我可以做这样的事情：

len_a = numpy.shape(a)[0]
sum_b = sum(b)
c = numpy.zeros((1, 0))
for idx in range(len_a):
    repeated_a = numpy.repeat(a[idx], b[idx])
    repeated_a = numpy.reshape(repeated_a, (1, numpy.shape(repeated_a)[0]))
    c = numpy.hstack((c, repeated_a))

但是，循环不是一个好的选择，因为它很慢。 我怎样才能让它更快？ 也许是某种形式的矢量化。

Answer 1

import numpy as np
import itertools

a = np.array([1, 2, 3 ,4, 5])
b = np.array([2, 3, 1, 2, 4])

list(itertools.chain(*[[i[0]] * i[1] for i in zip(a,b)]))

Answer 2

np.concatenate([np.repeat(a[i], b[i]) for i in range(len(a))])

输出：

array([1, 1, 2, 2, 2, 3, 4, 4, 5, 5, 5, 5])

Answer 3

您正在寻找为此目的制作的 built-in 重复功能。只需将两个数组都输入函数即可：

np.repeat(a,b)
#[1 1 2 2 2 3 4 4 5 5 5 5]

Answer 4

根据您对两个数组的了解，您可能会获得比本机提供的更多的性能：您知道这两个数组中可以出现什么样的数字吗？您知道生成的数组有多长，因此您可以利用它并预先初始化该内存吗？

为了说明这有什么不同，让我们看看您可以在没有更多技巧的情况下从 Numba 中挤出什么：

from numba import njit
import numpy as np

@njit
def f(a, b):
    s = b.sum()
    res = np.empty(s, dtype=np.int64)
    cs = 0
    for i in range(len(a)):
        bi = b[i]
        res[cs:cs+bi] = a[i]
        cs += bi
    return res

让我们使用 IPython 的魔法 %timeit 看看它与 np.repeat(a, b) 相比如何：

In [2]: a = np.random.randint(1, 6, 10000)
   ...: b = np.random.randint(1, 6, 10000)

In [3]: %timeit np.repeat(a, b)
135 µs ± 1.03 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [6]: %timeit f(a, b)
87 µs ± 485 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

这本身就是一个很好的改进。现在，如果您事先知道 b.sum() 是什么，就可以避免显式计算它，从而节省一些时间。或者，如果您知道 b 的值是如何分布的，则可以计算出数组大小的上限并在一天结束时截断 res[:cs]。

重复值 n 次，n 在一个数组中

Repeat values n many times, n being in an array

python

numpy

vectorization