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

创建数组的有效方法，该数组是numpy中可变长度范围的序列

2 个答案: