circular numpy array indices

2020-07-06 03:07发布

I have a 1-D numpy array a = [1,2,3,4,5,6] and a function that gets two inputs, starting_index and ending_index, and returns a[staring_index:ending_index].

Clearly I run into trouble when ending_index is smaller than starting_index. In this case, the function should start from the starting_index and traverse the vector a in a circular way, i.e., return all the elements coming after starting_index plus all the elements from index zero to ending_index.

For instance, if starting_index=4 and ending_index=1 then the output should be output = [5,6,1]. I can implement it with an if condition but I was wondering if there is any Pythonic and concise way to do it?

4条回答
混吃等死
2楼-- · 2020-07-06 03:27

Unfortunatly you cannot do this with slicing, you'll need to concatonate to two segments:

import numpy as np

a = [1, 2, 3, 4, 5, 6]
if starting_index > ending_index:
    part1 = a[start_index:]
    part2 = a[:end_index]
    result = np.concatenate([part1, part2])
else:
    result = a[start_index:end_index]
查看更多
祖国的老花朵
3楼-- · 2020-07-06 03:30

An alternative that you can use is the numpy roll function combined with indexing:

# -*- coding: utf-8 -*-
import numpy as np

def circular_array(starting_index, ending_index):

    idx = np.arange(1,7)
    idx = np.roll(idx, -starting_index)[:(len(idx)-starting_index+ending_index)%len(idx)]

    return idx


a = circular_array(4, 1)
print a
查看更多
我只想做你的唯一
4楼-- · 2020-07-06 03:34

This circles forever.

def circular_indices(lb, ub, thresh):
    indices = []
    while True:
        stop = min(ub, thresh)
        ix = np.arange(lb, stop)
        indices.append(ix)
        if stop != ub:
            diff = ub - stop
            lb = 0
            ub = diff
        else:
            break

    return np.concatenate(indices)
查看更多
姐就是有狂的资本
5楼-- · 2020-07-06 03:40

np.take has a wrap mode:

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

That's not quite what you want.

Actually, modulus does the same thing

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

But how about a small function that turns (4,1) into [4, 5, 6], or if you prefer [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))
查看更多
登录 后发表回答