在 numpy 中创建作为可变长度范围序列的数组的有效方法

Question

假设我有一个数组

import numpy as np
x=np.array([5,7,2])

我想创建一个数组，其中包含一系列范围堆叠在一起，每个范围的长度由 x 给出：

y=np.hstack([np.arange(1,n+1) for n in x])

有没有办法在没有列表理解或循环的速度损失的情况下做到这一点。（x 可能是一个非常大的数组）

结果应该是

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

Answer 1

您可以使用累积：

def my_sequences(x):
    x = x[x != 0] # you can skip this if you do not have 0s in x.

    # Create result array, filled with ones:
    y = np.cumsum(x, dtype=np.intp)
    a = np.ones(y[-1], dtype=np.intp)

    # Set all beginnings to - previous length:
    a[y[:-1]] -= x[:-1]

    # and just add it all up (btw. np.add.accumulate is equivalent):
    return np.cumsum(a, out=a) # here, in-place should be safe.

（一个警告：如果你的结果数组会更大，那么可能的大小np.iinfo(np.intp).max可能会因为运气不好而返回错误的结果而不是干净地出错......）

并且因为每个人都想要计时（与 Ophion 相比）方法：

In [11]: x = np.random.randint(0, 20, 1000000)

In [12]: %timeit ua,uind=np.unique(x,return_inverse=True);a=[np.arange(1,k+1) for k in ua];np.concatenate(np.take(a,uind))
1 loops, best of 3: 753 ms per loop

In [13]: %timeit my_sequences(x)
1 loops, best of 3: 191 ms per loop

当然，my_sequences当 的值x变大时，该函数不会表现不佳。

Answer 2

第一个想法；防止多次调用np.arange并且concatenate应该更快hstack：

import numpy as np
x=np.array([5,7,2])

>>>a=np.arange(1,x.max()+1)
>>> np.hstack([a[:k] for k in x])
array([1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 6, 7, 1, 2])

>>> np.concatenate([a[:k] for k in x])
array([1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 6, 7, 1, 2])

如果有许多非唯一值，这似乎更有效：

>>>ua,uind=np.unique(x,return_inverse=True)
>>>a=[np.arange(1,k+1) for k in ua]
>>>np.concatenate(np.take(a,uind))

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

您的案例的一些时间安排：

x=np.random.randint(0,20,1000000) 

原始代码

#Using hstack
%timeit np.hstack([np.arange(1,n+1) for n in x])
1 loops, best of 3: 7.46 s per loop

#Using concatenate
%timeit np.concatenate([np.arange(1,n+1) for n in x])
1 loops, best of 3: 5.27 s per loop

第一个代码：

#Using hstack
%timeit a=np.arange(1,x.max()+1);np.hstack([a[:k] for k in x])
1 loops, best of 3: 3.03 s per loop

#Using concatenate
%timeit a=np.arange(1,x.max()+1);np.concatenate([a[:k] for k in x])
10 loops, best of 3: 998 ms per loop

第二个代码：

%timeit ua,uind=np.unique(x,return_inverse=True);a=[np.arange(1,k+1) for k in ua];np.concatenate(np.take(a,uind))
10 loops, best of 3: 522 ms per loop

看起来最终代码的速度提高了 14 倍。

小健全性检查：

ua,uind=np.unique(x,return_inverse=True)
a=[np.arange(1,k+1) for k in ua]
out=np.concatenate(np.take(a,uind))

>>>out.shape
(9498409,)

>>>np.sum(x)
9498409