使用numpy重复arange函数

Question

我有一个整数值的数组。

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

我希望能够对a中的每个值应用一个arange函数，以便使其变成如下形式:

b = [0,1,0,0,1,2,3,1,2]
b "=" [arange(2),arange(1),arange(4),arange(0),arange(2)] 

实际上，我使用np.repeat函数根据数组a重复数组行，并希望有一个标记i将每个重复值链接到原始值，并具有识别号以区分它们。

我尝试使用np.vectorize但没有成功。

- Alexis

---

你可以展示更多的代码吗？你是怎么从 a 到 b 的呢？最重要的是，你尝试过什么？ - Inbar Rose

Answer 1

有更多“numpythonic”的做法。一种可能的方法是这样的：

import numpy as np
from numpy.lib.stride_tricks import as_strided

def concatenated_ranges(ranges_list) :
    ranges_list = np.array(ranges_list, copy=False)
    base_range = np.arange(ranges_list.max())
    base_range =  as_strided(base_range,
                             shape=ranges_list.shape + base_range.shape,
                             strides=(0,) + base_range.strides)
    return base_range[base_range < ranges_list[:, None]]

如果您只需要连接少数范围，那么Mr. E的纯Python解决方案可能是最佳选择，但是如果您需要连接至少100个范围，那么这种方法会更快。为了比较，我使用了从其他答案中提取出来的两个函数：

def junuxx(a) :
    b = np.array([], dtype=np.uint8)
    for x in a:
        b = np.append(b, np.arange(x))
    return b

def mr_e(a) :
    return reduce(lambda x, y: x + range(y), a, [])

以下是一些时间统计：

In [2]: a = [2, 1, 4, 0 ,2] # the OP's original example

In [3]: concatenated_ranges(a) # show it works!
Out[3]: array([0, 1, 0, 0, 1, 2, 3, 0, 1])

In [4]: %timeit concatenated_ranges(a)
10000 loops, best of 3: 31.6 us per loop

In [5]: %timeit junuxx(a)
10000 loops, best of 3: 34 us per loop

In [6]: %timeit mr_e(a)
100000 loops, best of 3: 2.58 us per loop

In [7]: a = np.random.randint(1, 10, size=(10,))

In [8]: %timeit concatenated_ranges(a)
10000 loops, best of 3: 27.1 us per loop

In [9]: %timeit junuxx(a)
10000 loops, best of 3: 79.8 us per loop

In [10]: %timeit mr_e(a)
100000 loops, best of 3: 7.82 us per loop

In [11]: a = np.random.randint(1, 10, size=(100,))

In [12]: %timeit concatenated_ranges(a)
10000 loops, best of 3: 57.4 us per loop

In [13]: %timeit junuxx(a)
1000 loops, best of 3: 756 us per loop

In [14]: %timeit mr_e(a)
10000 loops, best of 3: 149 us per loop

In [15]: a = np.random.randint(1, 10, size=(1000,))

In [16]: %timeit concatenated_ranges(a)
1000 loops, best of 3: 358 us per loop

In [17]: %timeit junuxx(a)
100 loops, best of 3: 9.38 ms per loop

In [18]: %timeit mr_e(a)
100 loops, best of 3: 8.93 ms per loop

- Jaime

---

很有趣的看到这些时间。谢谢。 - YXD

+1，我喜欢带有基准测试的答案！我得研究一下这个 as_strided 函数，它让我感到困惑。 - Junuxx

@Junuxx 这是一些小的numpy技巧，使其变得如此不可抗拒。来自numpy大师seberg的这些gists展示了它在工作中的一些很好的例子。 - Jaime

非常感谢。阅读完后，我在这里查看了一些文档（http://scipy-lectures.github.com/advanced/advanced_numpy/index.html）。它非常有帮助和实用。 - Alexis

Answer 2

我的回答与Junuxx的类似 - 我不确定你对b的回答是否是你想要的。

a = [2, 1, 4, 0 ,2]
reduce(lambda x, y: x+range(y), a, [])

给我。

[0, 1, 0, 0, 1, 2, 3, 0, 1]

- YXD

---

另一个不一致之处在于问题标题和标签提到了Numpy，但示例显示的是列表。您的代码可以在列表上运行，但不能在Numpy数组上运行（此外，在数组上不执行附加操作，而是进行向量加法）。 - Junuxx

这是正确的。我认为它提到了np.repeat，因为它与期望的行为相近，而不是特别针对numpy的问题。 - YXD

Answer 3

这段代码实现了你所描述的功能，即将a中的所有值连接起来形成一个aranges数组。尽管这意味着b中可能存在一些错误:

>>> a = [2, 1, 4, 0, 2]
>>> b = np.array([], dtype=np.uint8)
>>>for x in a:
>>>    b = np.append(b, np.arange(x))
>>> print b
array([0,1,0,0,1,2,3,0,1,])

在您解释需要一份排列列表之后，我认为以下方式更加有效：

>>> a = [2, 1, 4, 0, 2]
>>> b = [np.arange(x) for x in a]
>>> print b
[array([0, 1]), array([0]), array([0, 1, 2, 3]), array([], dtype=int32),
 array([0, 1])]

- Junuxx

---

当我开始学习Python的时候，拥有不同的解决方法总是很好的。这将有助于未来遇到的问题。我只是提到Junuxx的解决方案似乎比Mr. E的要快一些。 - Alexis

Answer 4

另一种更加内存高效的方法，通常稍微快一些：

import numpy as np
def concatenated_ranges2(ranges_list):
    cumsum = np.append(0, np.cumsum(ranges_list[:-1]))
    cumsum = np.repeat(cumsum, ranges_list)
    return np.arange(cumsum.shape[0]) - cumsum

测试这个函数和之前的一个：

>>> a = np.random.randint(1, 10, size=(1000,))
>>> %timeit concatenated_ranges(a)
10000 loops, best of 3: 142 us per loop
>>> %timeit concatenated_ranges2(a)
10000 loops, best of 3: 72.6 us per loop

- user2918340