从索引数组的函数返回等效的`:` [英] Return equivalent of `:` from function for indexing array

问题描述

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

I have a large array and a function that returns index lists into the array, i.e.,

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返回:.

A typical case is that the range is the entire array, range(n). Unfortunately, a[range(n)] scales with n while a[:] is of course constant-time. It's a pity that one cannot return : from get_idx.

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

What can I return from get_idx to use as an index for the entire array?

推荐答案

NumPy有一个帮助器

NumPy has a helper np.s_[] which can be used to construct slice and Ellipsis objects:

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> ]完全相同.

这篇关于从索引数组的函数返回等效的`:`的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！