Python/Numpy - 数组尾部切片环绕

6 投票
2 回答
2031 浏览
提问于 2025-04-16 15:51

我有两个一维数组,一个数组里有我感兴趣的一些值(叫它a),另一个数组提供了这些值在第一个数组中的位置(叫它b)。我知道b数组里的值总是递增的,除了某个地方(可能在任何位置),那里会出现一个下降的情况,因为它是从数组a的末尾回绕到开头。下面这个方法似乎可以用,但我觉得应该有更简单的方法。有没有人能给我推荐更好的办法?谢谢。

代码:

import numpy as np
a = np.arange(12)
b = np.array([5, 9, 2, 4])
#I want to generate these:
#[5,6,7,8,9]
#[9,10,11,0,1,2]
#[2,3,4]
#[4,5]

a = np.roll(a, -b[0], axis=0)
# Subtract off b[0] but ensure that all values are positive
b = (b-b[0]+len(a))%len(a)
for i, ind in enumerate(b):
   if i < len(b)-1:
      print a[b[i]:b[i+1]+1]
   else:
      print np.hstack((a[b[i]:len(a)], a[0]))

2 个回答

1

不确定这是否有帮助,但有一种快速的方法,不用去修改a的内存,方法是这样的:

import numpy as np

a = np.arange(12)
b = np.array([5, 9, 2, 4])
b = np.append(b, b[0])

b2 = b.copy()

b2[(np.diff(b)<0).nonzero()[0]+1:] += a.size

print [np.take(a, np.r_[b2[i]:b2[i+1]+1], mode='wrap') for i in range(b.size-1)]

print [np.roll(a, len(a)-b[i])[:b[i+1]-b[i]+1] for i in range(b.size-1)]

%timeit [np.take(a, np.r_[b2[i]:b2[i+1]+1], mode='wrap') for i in range(b.size-1)]
# 10000 loops, best of 3: 28.6 µs per loop

%timeit [np.roll(a, len(a)-b[i])[:b[i+1]-b[i]+1] for i in range(b.size-1)]
# 10000 loops, best of 3: 77.7 µs per loop
3

稍微简短了一点,不过我想我还能做得更好……

import numpy as np

a = np.arange(12)
b = np.array([5, 9, 2, 4])
b = np.append(b, b[0])

for i in range(0, len(b)-1):
    print np.roll(a, len(a)-b[i])[:b[i+1]-b[i]+1]

撰写回答