从不同大小的小数组构建单个numpy数组

Question

我有一个值的数组 x。给定一个'开始'和'结束'的索引，我需要用 x 的子数组来构建一个数组 y。

import numpy as np
x = np.arange(20)
start = np.array([2, 8, 15])
stop = np.array([5, 10, 20])
nsubarray = len(start)

我希望 y 是这样的：

y = array([ 2,  3,  4,  8,  9, 15, 16, 17, 18, 19])

（实际上我使用的数组要大得多）。

构建 y 的一种方法是使用列表推导式，但之后需要把列表扁平化：

import itertools as it
y = [x[start[i]:stop[i]] for i in range(nsubarray)]
y = np.fromiter(it.chain.from_iterable(y), dtype=int)

我发现其实用 for 循环会更快：

y = np.empty(sum(stop - start), dtype = int)
a = 0
for i in range(nsubarray):
    b = a + stop[i] - start[i]
    y[a:b] = x[start[i]:stop[i]]
    a = b

我在想有没有人知道我可以怎么优化这个？非常感谢！

编辑

以下是所有测试的时间：

import numpy as np
import numpy.random as rd
import itertools as it


def get_chunks(arr, start, stop):
    rng = stop - start
    rng = rng[rng!=0]      #Need to add this in case of zero sized ranges
    np.cumsum(rng, out=rng)
    inds = np.ones(rng[-1], dtype=np.int)
    inds[rng[:-1]] = start[1:]-stop[:-1]+1
    inds[0] = start[0]
    np.cumsum(inds, out=inds)
    return np.take(arr, inds)


def for_loop(arr, start, stop):
    y = np.empty(sum(stop - start), dtype = int)
    a = 0
    for i in range(nsubarray):
        b = a + stop[i] - start[i]
        y[a:b] = arr[start[i]:stop[i]]
        a = b
    return y

xmax = 1E6
nsubarray = 100000
x = np.arange(xmax)
start = rd.randint(0, xmax - 10, nsubarray)
stop = start + 10

结果是：

In [379]: %timeit np.hstack([x[i:j] for i,j in it.izip(start, stop)])
1 loops, best of 3: 410 ms per loop

In [380]: %timeit for_loop(x, start, stop)
1 loops, best of 3: 281 ms per loop

In [381]: %timeit np.concatenate([x[i:j] for i,j in it.izip(start, stop)])
10 loops, best of 3: 97.8 ms per loop

In [382]: %timeit get_chunks(x, start, stop)
100 loops, best of 3: 16.6 ms per loop

Answer 1

使用切片代替np.arrays可以吗？

import numpy as np
x = np.arange(10)
start = slice(2, 8)
stop = slice(5, 10)

print np.concatenate((x[start], x[stop]))

Answer 2

这对我来说几乎快了三倍，几乎所有的时间差都来自于把 fromiter 替换成 concatenate：

import numpy as np
from itertools import izip

y = [x[b:e] for b, e in izip(start, stop)]
y = np.concatenate(y)

Answer 3

也许可以用 zip、np.arange 和 np.hstack 来实现：

np.hstack([np.arange(i, j) for i,j in zip(start, stop)])

Answer 4

这有点复杂，但速度很快。基本上，我们的做法是根据向量加法来创建索引列表，并使用 np.take，而不是用任何Python的循环：

def get_chunks(arr, start, stop):
     rng = stop - start
     rng = rng[rng!=0]      #Need to add this in case of zero sized ranges
     np.cumsum(rng, out=rng)
     inds = np.ones(rng[-1], dtype=np.int)
     inds[rng[:-1]] = start[1:]-stop[:-1]+1
     inds[0] = start[0]
     np.cumsum(inds, out=inds)
     return np.take(arr, inds)

检查一下它是否返回了正确的结果：

xmax = 1E6
nsubarray = 100000
x = np.arange(xmax)
start = np.random.randint(0, xmax - 10, nsubarray)
stop = start + np.random.randint(1, 10, nsubarray)

old = np.concatenate([x[b:e] for b, e in izip(start, stop)])
new = get_chunks(x, start, stop)
np.allclose(old,new)
True

一些时间测试：

%timeit np.hstack([x[i:j] for i,j in zip(start, stop)])
1 loops, best of 3: 354 ms per loop

%timeit np.concatenate([x[b:e] for b, e in izip(start, stop)])
10 loops, best of 3: 119 ms per loop

%timeit get_chunks(x, start, stop)
100 loops, best of 3: 7.59 ms per loop