How to find the index of an element in numpy?

In NumPy, you can use the np.where() function to find the index of elements. The np.where() function returns the index of elements that meet the specified condition. Below is an example of using the np.where() function to find element indexes.

import numpy as np

arr = np.array([1, 2, 3, 4, 5, 6])

# 查找元素3的索引
index = np.where(arr == 3)
print(index)  # 输出 (array([2]),)

# 查找元素大于4的索引
index = np.where(arr > 4)
print(index)  # 输出 (array([4, 5]),)

# 查找元素不等于2的索引
index = np.where(arr != 2)
print(index)  # 输出 (array([0, 2, 3, 4, 5]),)

Please note that the function np.where() returns a tuple containing the index arrays of elements that satisfy the condition. To retrieve the index array, you can use index[0].