Python/Numpy - Wrap Slice Around End of Array
I have two 1D arrays, one that has some values of interest (a) and another that provides indices into that array (b). I know that the values in b always increase, except at one point (could be anywhere) where the number decreases since it rolls from the end to the beginning of array a. The method below seems to work, but I just think that a cleaner way must exist. Can anyone suggest something better? Thanks.
Code:
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 en开发者_开发知识库sure 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]))
A bit shorter, but maybe I can still do better...
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]
Not sure if this helps, but a fast way without messing with the memory of a
would be this:
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
精彩评论