python - 从索引数组的函数返回等价于 `:`

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)] 与 n 成比例，而 a[:] 当然是常数时间。很遗憾，不能从 get_idx 返回 :。

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

Answer 1

看看slice

def get_x():
    return slice(2)

a=list(range(100))
a[get_x()]

将返回[0, 1]

更新

根据您的需要，get_x 函数应该是

def get_x(k, n):
    return slice(n if k > 6 else k)

更新

正如@Eric 正确指出的那样，最好传递 None 而不是 n。 所以功能将是:

def get_x(k):
    return slice(None if k > 6 else k)