Python中numpy库的where函数_numpy.where

numpy.where() 的两种用法：

1、numpy.where(condition, x, y)
满足条件(condition)，输出x，不满足输出y。

import numpy as np

A = np.arange(-5, 10)
print(A)
#[-5 -4 -3 -2 -1  0  1  2  3  4  5  6  7  8  9]

B = np.where(A, 1, -1)  #0为False，所以0变为-1
print(B)
#[ 1  1  1  1  1 -1  1  1  1  1  1  1  1  1  1]

C = np.where(A > 2, 1, -1)
print(C)
#[-1 -1 -1 -1 -1 -1 -1 -1  1  1  1  1  1  1  1]

np.random.seed(1)
arr = np.random.randn(4, 4)  #正态分布的数据
print(arr)
'''
[[ 1.62434536 -0.61175641 -0.52817175 -1.07296862]
 [ 0.86540763 -2.3015387   1.74481176 -0.7612069 ]
 [ 0.3190391  -0.24937038  1.46210794 -2.06014071]
 [-0.3224172  -0.38405435  1.13376944 -1.09989127]]
'''

D = np.where(arr > 0, 1, -1)  #将所有正值替换为1，所有负值替换为-1
print(D)
'''
[[ 1 -1 -1 -1]
 [ 1 -1  1 -1]
 [ 1 -1  1 -1]
 [-1 -1  1 -1]]
'''

E = np.where(arr > 0, 1, arr)  #只将所有正值替换为1
print(E)
'''
[[ 1.         -0.61175641 -0.52817175 -1.07296862]
 [ 1.         -2.3015387   1.         -0.7612069 ]
 [ 1.         -0.24937038  1.         -2.06014071]
 [-0.3224172  -0.38405435  1.         -1.09989127]]
'''

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
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

2、numpy.where(condition)
只有条件 (condition)，没有x和y，输出满足条件 (即非0) 元素的索引 (等价于numpy.nonzero)。这里的索引以tuple的形式给出，通常原数组有多少维，输出的tuple中就包含几个数组，分别对应符合条件元素的各维索引。

import numpy as np

A = np.array([1, 3, 5, 7, 9, 11])

B = np.where(A>3)  #返回索引
print(B)
#(array([2, 3, 4, 5]),)

C = A[np.where(A>3)]  #等价于A[A>3]
print(C)
#[ 5  7  9 11]

1
2
3
4
5
6
7
8
9
10
11
12