numpy数组中所有元素的范围

Range of all elements in numpy array

假设我有一个 1d 数组 a,其中每个元素我希望有一个范围,其大小存储在 ranges:

a = np.array([10,9,12])
ranges = np.array([2,4,3])

所需的输出将是:

np.array([10,11,9,10,11,12,12,13,14])

我当然可以使用 for 循环,但我更喜欢完全矢量化的方法。 np.repeat 允许通过设置 repeats= 多次重复 a 中的元素,但我不知道类似的 numpy 函数特别处理上述问题。

试试这个:

a = np.array([10,9,12])
ranges = np.array([2,4,3])
[i for j in range(len(a)) for i in range(a[j],a[j]+ranges[j])]
# [10, 11,  9, 10, 11, 12, 12, 13, 14]

如果numpy.array试试这个:

np.array([i for j in range(len(a)) for i in range(a[j],a[j]+ranges[j])])
# array([10, 11,  9, 10, 11, 12, 12, 13, 14])

短数组的运行时间:

使用 pandas 可能会更容易:

>>> import pandas as pd
>>> x = pd.Series(np.repeat(a, ranges))
>>> x + x.groupby(x).cumcount()
0    10
1    11
2     9
3    10
4    11
5    12
6    12
7    13
8    14
dtype: int64
>>> 

如果你想要一个 numpy 数组:

>>> x.add(x.groupby(x).cumcount()).to_numpy()
array([10, 11,  9, 10, 11, 12, 12, 13, 14], dtype=int64)
>>> 
>>> np.hstack([np.arange(start, start+size) for start, size in zip(a, ranges)])
array([10, 11,  9, 10, 11, 12, 12, 13, 14])

有人问时间问题,所以我使用 Jupyter notebook 单元格中的 %timeit 魔术功能,以非常简单的方式比较了三种解决方案(到目前为止)的时间。

我设置如下:

N = 1
a = np.array([10,9,12])
a = np.tile(a, N)
ranges = np.array([2,4,3])
ranges = np.tile(ranges, N)
a.shape, ranges.shape

所以我可以轻松扩展(尽管事情不是 运行dom,而是重复)。

那我运行:

%timeit np.hstack([np.arange(start, start+size) for start, size in zip(a, ranges)])

,

%timeit x = pd.Series(np.repeat(a, ranges)); x.add(x.groupby(x).cumcount()).to_numpy()

%timeit np.array([i for j in range(len(a)) for i in range(a[j],a[j]+ranges[j])])

结果如下: N = 1:

  1. 每个循环 9.81 µs ± 481 ns(7 次运行的平均值 ± 标准偏差,每次 100000 次循环)
  2. 每个循环 568 µs ± 20.8 µs(7 次运行的平均值 ± 标准偏差,每次 1000 次循环)
  3. 每个循环 3.53 µs ± 81.4 ns(7 次运行的平均值 ± 标准偏差,每次 100000 次循环)

N = 10:

  1. 每个循环 63.4 µs ± 976 ns(7 次运行的平均值 ± 标准偏差,每次 10000 次循环)
  2. 每个循环 575 µs ± 15.5 µs(7 次运行的平均值 ± 标准偏差,每次 1000 次循环)
  3. 每个循环 25.1 µs ± 698 ns(7 次运行的平均值 ± 标准偏差,每次 10000 次循环)

N = 100:

  1. 每个循环 612 µs ± 12.4 µs(7 次运行的平均值 ± 标准偏差,每次 1000 次循环)
  2. 每个循环 608 µs ± 25.4 µs(7 次运行的平均值 ± 标准偏差,每次 1000 次循环)
  3. 每个循环 237 µs ± 9.62 µs(7 次运行的平均值 ± 标准偏差,每次 1000 次循环)

N = 1000:

  1. 每个循环 6.09 ms ± 52 µs(7 次运行的平均值 ± 标准偏差,每次 100 次循环)
  2. 每个循环 852 µs ± 2.66 µs(7 次运行的平均值 ± 标准偏差,每次 1000 次循环)
  3. 每个循环 2.44 ms ± 43.3 µs(7 次运行的平均值 ± 标准偏差,每次 100 次循环)

因此,当事情达到 1000 个或更多元素的数组时,Pandas 解决方案获胜,但 Python 双列表理解在这一点之前表现出色。 np.hstack 可能因为额外的内存分配和复制而失败,但这是一个猜测。另请注意,对于每个数组大小,Pandas 解决方案的时间几乎相同。

警告仍然存在,因为有重复的数字,并且所有值都是相对较小的整数。这真的无关紧要,但我(还)没有打赌。 (例如,Pandas groupby 功能可能会因为重复的数字而很快。)


奖励:OP 在评论中声明“现实生活中的数组大约有 1000 个元素,但是 运行ges 运行 从 100 到 1000。所以变得相当大 – pr94 ”。 所以我将计时测试调整为以下内容:

import numpy as np
import pandas as pd

N = 1000
a = np.random.randint(100, 1000, N)
# This is how I understand "ranges ranging from 100 to 1000"
ranges = np.random.randint(100, 1000, N)

%timeit np.hstack([np.arange(start, start+size) for start, size in zip(a, ranges)])
%timeit  x = pd.Series(np.repeat(a, ranges)); x.add(x.groupby(x).cumcount()).to_numpy()
%timeit np.array([i for j in range(len(a)) for i in range(a[j],a[j]+ranges[j])])

结果为:

  • hstack:每个循环 2.78 ms ± 38.6 µs(7 次运行的平均值 ± 标准偏差,每次 100 次循环)
  • pandas:每个循环 18.4 ms ± 663 µs(7 次运行的平均值 ± 标准偏差,每次 100 次循环)
  • 双重列表理解:每个循环 64.1 毫秒 ± 427 微秒(7 次运行的平均值 ± 标准差,每次 10 次循环)

这表明我提到的那些注意事项(至少以某种形式)似乎确实存在。但是人们应该仔细检查这个测试代码是否实际上是最相关和最合适的,以及它是否正确。

使用 Numba 编译函数可能会更快地解决此问题:

@nb.jit 
def expand_range(values, counts): 
    n = len(values) 
    m = np.sum(counts) 
    r = np.zeros((m,), dtype=values.dtype) 
    k = 0 
    for i in range(n): 
        x = values[i] 
        for j in range(counts[i]): 
            r[k] = x + j 
            k += 1 
    return r

在非常小的输入上:

%timeit expand_range(a, ranges)                                                                                                                                                
# 1.16 µs ± 126 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%timeit x = pd.Series(np.repeat(a, ranges)); x.add(x.groupby(x).cumcount()).to_numpy()                                                                                         
# 617 µs ± 4.32 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.hstack([np.arange(start, start+size) for start, size in zip(a, ranges)])                                                                                            
# 25 µs ± 2.2 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit np.array([i for j in range(len(a)) for i in range(a[j],a[j]+ranges[j])])                                                                                               
# 13.5 µs ± 929 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

以及更大的输入:

b = np.random.randint(0, 1000, 1000)
b_ranges = np.random.randint(1, 10, 1000)

%timeit expand_range(b, b_ranges)                                                                                                                                              
# 5.07 µs ± 98.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit x = pd.Series(np.repeat(a, ranges)); x.add(x.groupby(x).cumcount()).to_numpy()                                                                                         
# 617 µs ± 4.32 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.hstack([np.arange(start, start+size) for start, size in zip(a, ranges)])                                                                                            
# 25 µs ± 2.2 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit np.array([i for j in range(len(a)) for i in range(a[j],a[j]+ranges[j])])                                                                                               
# 13.5 µs ± 929 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

这些表明,使用基于 Numba 的方法获胜,速度增益至少是目前提出的任何其他方法的 100 倍。


数字更接近 OP 的其中一条评论中指出的数字:

b = np.random.randint(10, 1000, 1000)
b_ranges = np.random.randint(100, 1000, 1000)

%timeit expand_range(b, b_ranges)                                                                                                                                              
# 1.5 ms ± 67.9 µs per loop (mean ± std. dev. of 7 runs, 1000

%timeit x = pd.Series(np.repeat(b, b_ranges)); x.add(x.groupby(x).cumcount()).to_numpy()                                                                                       
# 91.8 ms ± 6.53 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit np.hstack([np.arange(start, start+size) for start, size in zip(b, b_ranges)])                                                                                          
# 10.7 ms ± 402 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit np.array([i for j in range(len(b)) for i in range(b[j],b[j]+b_ranges[j])])                                                                                             
# 144 ms ± 4.54 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

这至少比其他的高出 7 倍。