Python 如何在numpy中高效地连接多个arange调用？

我想在

cnt

值的向量上对调用进行向量化，如

numpy.arange（0，cnt_I）

，并像以下代码片段那样连接结果：

import numpy
cnts = [1,2,3]
numpy.concatenate([numpy.arange(cnt) for cnt in cnts])

array([0, 0, 1, 0, 1, 2])

不幸的是，由于临时数组和列表理解循环，上面的代码内存效率非常低

---

在numpy中有没有一种方法可以更有效地做到这一点？

np。tril_Index

几乎可以为您做到这一点：

In [28]: def f(c):
   ....:     return np.tril_indices(c, -1)[1]

In [29]: f(10)
Out[29]:
array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1,
       2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 8])

In [33]: %timeit multirange(range(10))
10000 loops, best of 3: 93.2 us per loop

In [34]: %timeit f(10)
10000 loops, best of 3: 68.5 us per loop

当维度较小时，比@Warren Weckesser

多量程

快得多

但是当维度变大时，速度会变慢（@hpaulj，你有一个很好的观点）：

---

尝试以下方法解决内存问题，效率几乎相同

out = np.empty((sum(cnts)))
k = 0
for cnt in cnts:
    out[k:k+cnt] = np.arange(cnt)
    k += cnt

---

因此不使用串联。

这里有一个完全矢量化的函数：

def multirange(counts):
    counts = np.asarray(counts)
    # Remove the following line if counts is always strictly positive.
    counts = counts[counts != 0]

    counts1 = counts[:-1]
    reset_index = np.cumsum(counts1)

    incr = np.ones(counts.sum(), dtype=int)
    incr[0] = 0
    incr[reset_index] = 1 - counts1

    # Reuse the incr array for the final result.
    incr.cumsum(out=incr)
    return incr

def multirange_original(counts):
    ranges = np.concatenate([np.arange(count) for count in counts])
    return ranges

这里是@Developer答案的一个变体，它只调用

arange

一次：

def multirange_loop(counts):
    counts = np.asarray(counts)
    ranges = np.empty(counts.sum(), dtype=int)
    seq = np.arange(counts.max())
    starts = np.zeros(len(counts), dtype=int)
    starts[1:] = np.cumsum(counts[:-1])
    for start, count in zip(starts, counts):
        ranges[start:start + count] = seq[:count]
    return ranges

这是原始版本，作为函数编写：

def multirange(counts):
    counts = np.asarray(counts)
    # Remove the following line if counts is always strictly positive.
    counts = counts[counts != 0]

    counts1 = counts[:-1]
    reset_index = np.cumsum(counts1)

    incr = np.ones(counts.sum(), dtype=int)
    incr[0] = 0
    incr[reset_index] = 1 - counts1

    # Reuse the incr array for the final result.
    incr.cumsum(out=incr)
    return incr

def multirange_original(counts):
    ranges = np.concatenate([np.arange(count) for count in counts])
    return ranges

演示：

使用更大的计数数组比较计时：

In [299]: counts = np.random.randint(1, 50, size=50)

In [300]: %timeit multirange_original(counts)
10000 loops, best of 3: 114 µs per loop

In [301]: %timeit multirange_loop(counts)
10000 loops, best of 3: 76.2 µs per loop

In [302]: %timeit multirange(counts)
10000 loops, best of 3: 26.4 µs per loop

---

这正是我想要的，只是我真的想避免Python循环+1但等待更长的时间，看看是否有真正的矢量化解决方案。不幸的是，嵌套的Python speed循环将使这一过程变得相当缓慢，尽管它在避免内存分配方面做得很好。破坏性的累积和（覆盖自身）是一个好主意，谢谢每循环114µs”-减少时间x5可能是一个很好的练习，但它实用吗？什么是现实世界的维度首先引发了这个问题？原始代码更容易理解。当我将

50

增加到

5000

时，

multirange

比原始代码慢一点。@hpaulj:+1作为您的第一条评论。这正好为约瑟夫提供了另一个选择，并没有解决为什么一开始就需要这样一个函数的问题。此外，如果速度和内存使用都非常重要，我可能会在Cython中实现该功能。