如何在NumPy数组中获取索引值而非元素值

现有如下Python NumPy代码，当前运行输出为数组元素值[20 25 30 35 40 45 50 55]，需求改为获取对应索引值，期望输出为[0,5,10,15,20,25,30,35]：

import numpy as np

data=np.arange(20,60)
def number_of_dies_probed(inp_arr)-> np.array:

    inp_arr=inp_arr.flatten()
    enter_numbers=int(input('Enter number of dices to be probed:'))#for example input value 8
    vb=((len(inp_arr))/enter_numbers)

    dies_to_probe=(inp_arr[np.arange(0,inp_arr.shape[0],int(vb))])
    return dies_to_probe
mn=number_of_dies_probed(data)
print(mn)

修改方案 #

核心调整是直接返回生成的索引数组，而非原数组的元素值，具体改动如下：

修改后的完整代码 #

import numpy as np

data=np.arange(20,60)
def number_of_dies_probed(inp_arr)-> np.array:
    inp_arr=inp_arr.flatten()
    enter_numbers=int(input('Enter number of dices to be probed:'))#for example input value 8
    vb=(len(inp_arr))/enter_numbers
    # 直接返回索引序列，不再取原数组元素
    dies_to_probe=np.arange(0,inp_arr.shape[0],int(vb))
    return dies_to_probe
mn=number_of_dies_probed(data)
print(mn)

改动说明 #

• 移除原代码中inp_arr[np.arange(...)]的元素取值操作，直接将np.arange生成的索引数组赋值给dies_to_probe；
• 当输入8时，vb=40/8=5，np.arange(0,40,5)会生成[0 5 10 15 20 25 30 35]，完全匹配期望输出。

内容的提问来源于stack exchange，提问作者Bhar Jay