[Python] np.where() - 조건에 맞는 인덱스 찾기

오늘은 넘파이 어레이에 조건에 맞는 인덱스를 찾는 함수에 포스팅하겠다.

 

np.where() - 1D

[ 1D array 생성 ]

우선 0~100까지의 어레이를 구성하겠다. 사이즈는 10.

import numpy as np

random_num = np.random.randint(0,100,size=10)
print(random_num)
# array([84, 63, 80, 84, 90, 98, 15, 92,  5, 12])

 

[최댓값, 최솟값 출력]

예시로 넘파이 어레이의 최댓값과 최솟값의 인덱스를 찾기 위해 각 값들을 출력해보겠다.

print(random_num.max(), random_num.min())
# 98 5

 

[np.where 사용]

이후 np.where()를 통해 각 값들이 random_num과 동일한 값이라면 해당 인덱스를 출력하도록 한다.

np.where(random_num == random_num.max())
# (array([5], dtype=int64),)

np.where(random_num == random_num.min())
# (array([8], dtype=int64),)

다음과 같이 최댓값, 최솟값이 위치해있는 인덱스가 출력된다.

 

뿐만 아니라 조건에 해당하는 값이 두 개 이상이더라도 모두 출력이 된다.

np.where(random_num < random_num.max())
# (array([0, 1, 2, 3, 4, 6, 7, 8, 9], dtype=int64),)

 

np.where() - 2차원 이상

[2D array 생성]

2차원 이상의 경우도 마찬가지로  np.where()은 동일하게 사용이 된다.

random_num_2D = np.random.randint(0,100,size=[3,3])
print(random_num_2D)
"""
[[ 9 85 99]
 [85 49  7]
 [ 8 96 14]]
"""

 

[최댓값 출력]

print(random_num_2D.max())
# 99

 

[최댓값 인덱스 찾기]

np.where(random_num_2D == random_num_2D.max())
# (array([0], dtype=int64), array([2], dtype=int64))

 

np.where() - if와 같은 역할

np.where()은 if와 같이 조건부 역할을 할 수도 있다.

 

[최댓값은 100, 나머지는 0]

np.where(random_num_2D == random_num_2D.max(), 100, 0)
"""
array([[  0,   0, 100],
       [  0,   0,   0],
       [  0,   0,   0]])
"""

 

[50 이상는 1, 이하는 0]

np.where(random_num_2D < 50, 1, 0)
"""
array([[1, 0, 0],
       [0, 1, 1],
       [1, 0, 1]])
"""

 

다음과 같이 np.where()을 잘 활용한다면 조건에 맞는 인덱스를 찾을 수 있을 뿐만 아니라 해당 조건에 부합하면 원하는 값으로 변경도 가능하다.