从索引数组的函数返回等效的`：`

Question

我有一个大数组和一个将索引列表返回到数组的函数，即

import numpy

n = 500
a = numpy.random.rand(n)

def get_idx(k):
    # More complicated in reality
    return range(n) if k > 6 else range(k)

data = a[get_idx(29)]
data = a[get_idx(30)]
# ...

典型的情况是范围是整个数组range(n)。不幸的是，a[range(n)]使用nwhile进行缩放a[:]当然是固定时间的。可惜一个人不能:从那里回来get_idx。

我可以从中返回什么get_idx用作整个数组的索引？

Answer 1

NumPy有一个np.s_[]可用于构造slice和Ellipsis对象的助手：

def get_idx(k):
    return np.s_[:] if k > 6 else np.s_[:k]

    # or even np.s_[:None if k > 6 else k]

一般而言，a[np.s_[ <stuff> ]]与完全相同a[ <stuff> ]。