np.where查找索引

Posted on | In , ,

np.where

np.where(condition, x, y)

满足条件(condition)，输出x，不满足输出y。

1
2
3
np.where([[True,False], [True,True]],    # 官网上的例子
[[1,2], [3,4]],
[[9,8], [7,6]])

输出

1
2
array([[1, 8],
[3, 4]])

上面这个例子的条件为[[True,False], [True,False]]，分别对应最后输出结果的四个值。第一个值从[1,9]中选，因为条件为True，所以是选1。第二个值从[2,8]中选，因为条件为False，所以选8，后面以此类推
这里的true指的就是选前面的，false就是指选后面的

1
2
3
4
5
6
7
>>> a = 10
>>> np.where([[a > 5,a < 5], [a == 10,a == 7]],
[["chosen","not chosen"], ["chosen","not chosen"]],
[["not chosen","chosen"], ["not chosen","chosen"]])

array([['chosen', 'chosen'],
['chosen', 'chosen']], dtype='<U10')

np.where(condition)

只有条件 (condition)，没有x和y，则输出满足条件 (即非0) 元素的坐标（注意这里返回的是坐标）

1
2
3
4
5
>>> a = np.array([2,4,6,8,10])
>>> np.where(a > 5)             # 返回索引
(array([2, 3, 4]),)   
>>> a[np.where(a > 5)]              # 等价于 a[a>5]
array([ 6,  8, 10])

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
>>> a = np.arange(27).reshape(3,3,3)
>>> a
array([[[ 0,  1,  2],
[ 3,  4,  5],
[ 6,  7,  8]],

[[ 9, 10, 11],
[12, 13, 14],
[15, 16, 17]],

[[18, 19, 20],
[21, 22, 23],
[24, 25, 26]]])

>>> np.where(a > 5)
(array([0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2]),
array([2, 2, 2, 0, 0, 0, 1, 1, 1, 2, 2, 2, 0, 0, 0, 1, 1, 1, 2, 2, 2]),
array([0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2]))

注意这里的最终结果的坐标是要竖着看的，即（0，2，0），（0，2，1）….

  • 这个方法只能用在array上面，如果需要list的话需要np.asarray
1
2
3
4
5
6
7
8
9
10
import numpy as np

ids = [[3], [34], [5]]
ids_o = [[3]]

remove_list = filter(lambda i: i not in ids_o, ids)
print(np.asarray(ids))
for i in remove_list:
index = np.where(np.asarray(ids) == i)
print(index)

结果

1
2
3
4
5
6
[[ 3]
[34]
[ 5]]
(array([1]), array([0]))
(array([2]), array([0]))
[Finished in 0.2s]