python - 循环 numpy 数组索引

Question

我有一个一维 numpy 数组 a = [1,2,3,4,5,6] 和一个获取两个输入的函数，starting_index 和ending_index，并返回a[staring_index:ending_index]。

显然，当 ending_index 小于 starting_index 时，我遇到了麻烦。在这种情况下，函数应该从 starting_index 开始，循环遍历向量 a，即返回 starting_index 之后的所有元素加上 index 中的所有元素零到 ending_index。

例如，如果 starting_index=4 和 ending_index=1 那么输出应该是 output = [5,6,1]。我可以用 if 条件来实现它，但我想知道是否有任何 Pythonic 和简洁的方法来实现它？

Answer 1

np.take 有一个wrap 模式:

In [171]: np.take(np.arange(1,7),range(4,7),mode='wrap')
Out[171]: array([5, 6, 1])

这不是你想要的。

实际上，模数做同样的事情

In [177]: a[np.array([4,5,6])%6]
Out[177]: array([5, 6, 1])

但是，将 (4,1) 转换为 [4, 5, 6] 的小函数怎么样，或者如果您更喜欢 [4, 5 , 0]?

def foo(a, start, stop): 
    # fn to convert your start stop to a wrapped range
    if stop<=start:
        stop += len(a)
    return np.arange(start, stop)%len(a)

a[foo(a,4,1)]  # or
np.take(a,foo(a,4,1))