python numpy的where函数

where函数:快速查找数组中大于某个值的索引```

import numpy as np
arr = np.array(
[[1,2,3],
[2,3,4]]
)
print(arr.shape) #(2,3)
t_index = np.where(arr>2)
print(t_index)
#(array([0, 1, 1], dtype=int64), array([2, 1, 2], dtype=int64))
#表示坐标为（0，2）（1，1）（1，2），即3，3，4三个数大于2

arr1 = arr[t_index[0],t_index[1]]
print('arr1',arr1)

其输出为：

arr1 [3 3 4]

如果只想要其中某一个值（比如（0，2），第一行第三列）：

arr2 = arr[0,2]
print('arr2',arr2) #输出为3，得到的是一个数

用一个复杂一些的数组，shape为（1，4，3，3）

arr1 = np.array(
[[[[ 0.17310286, -0.00099834, -0.01867144],
[ 0.26144078, 0.06111889, -0.05025296],
[ 0.16852978, 0.12282149, 0.03148762]],

[[ 0.12197495, 0.07626072, -0.03471078],
[ 0.10431765, 0.08020866, 0.02596528],
[-0.01284805, 0.02070661, -0.04171723]],

[[-0.08971099, -0.25008383, -0.11431476],
[-0.15451074, -0.3139743, -0.31134847],
[-0.20343271, -0.32051, -0.2928716 ]],

[[ 0.08037898, 0.02238214, 0.09162696],
[-0.10115768, -0.07861736, 0.00288352],
[-0.06310503, -0.15695804, -0.13134862]]]]
)

用其（0，0，3，3）的数据，
像上面的例子一样取数组中大于0.6的值，

t_index = np.where(arr>0.6)
print('t_index',t_index)

其输出为：

t_index (array([0, 1, 1, 2, 2], dtype=int64), array([0, 0, 1, 0, 1], dtype=int64))

取其中的（1，0）和（0，2），即第二行第一列和第一行第三列，输出会变成一维数组的形式（数组切片）

arr2 = arr1[0, 0, [1, 0], [0,2]]
print('arr2',arr2)

输出为：

arr2 [ 0.26144078 -0.01867144]

如果写成arr2 = arr1[0, 0, [[1, 0], [0,2]]]，得到的是：

[[[ 0.26144078  0.06111889 -0.05025296],   
[ 0.17310286 -0.00099834 -0.01867144]],
[[ 0.17310286 -0.00099834 -0.01867144],
  [ 0.16852978  0.12282149  0.03148762]]] 

可以看到，原来（0，0，3，3）中（3，3）部分的第一行和第二行对换，成为第一部分，第三行和第一行对换，成为第二部分，且每一行都完整保留