将计数列表转换为数字列表的最有效方法

Question

我有一个计数列表，其中每个索引代表一个数字，它的计数代表该数字在列表中的数量：

a = [3,5,1,2]

变成

b = [0,0,0,1,1,1,1,1,2,3,3]

我在想我们可以做这样的事情：

b = []
for ix, el in enumerate(a):
    b.extend([ix]*a[ix])
print(b)

但是，如果我没有记错的话，将它放入列表 b 需要 k (count val) 时间，因为扩展需要 k 时间，但它也必须完成 n 次，从而为我们提供 n*k 的运行时间，其中 n 是indicies 和 k 是每个索引的计数数

另一个想法是，我们可以拥有一个纯元素数组，而不是拥有一个计数数组：

a = [[0,0,0],[1,1,1,1,1],[2],[3,3]]

但要弄平它仍然需要相当长的时间（我相信 n*k 时间）

b = [item for sublist in a for item in sublist]

有没有办法使这更有效？ 也许转换为字符串删除所有括号并转换回列表？

Answer 1

您可以使用 numpy 的np.repeat来实现高性能方法：

np.repeat(np.arange(len(a)), a)
# array([0, 0, 0, 1, 1, 1, 1, 1, 2, 3, 3])

---

以下是时间——

a_large = np.concatenate([a]*10_000, axis=0)

def op(a):
    b = []
    for ix, el in enumerate(a):
        b.extend([ix]*a[ix])

def yatu(a):
    np.repeat(np.arange(len(a)), a)

%timeit op(a_large)
# 17.1 ms ± 422 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit yatu(a_large)
# 368 µs ± 1.91 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)