[Numpy] np.where 사용법 : 조건 만족하는 위치 인덱스 찾기

파이썬 넘파이 np.where : 조건 만족 인덱스 반환

파이썬에서 넘파이를 사용하다보면, np.where 함수를 자주 만나게 됩니다.

만날때 마다 복잡해보이지만, 사실 알고보면 매우 간단한 np.where 함수의 사용 방법을 간단히 살펴보겠습니다.

 

np.where 기본 사용법 (인덱스만 반환)

우선, 기본적으로 넘파이의 np.where 함수를 사용하는 방법은 np.where 내에 조건만 적는 방법입니다.

numpy array를 1개 선언해본 뒤, np.where로 조건을 만족하는 위치의 인덱스를 가져오는 예시를 보여드리겠습니다.

import numpy as np

a = np.arange(5, 15) # array([ 5,  6,  7,  8,  9, 10, 11, 12, 13, 14])

np.where(a > 10) # 6 ~ 9번 위치가 만족 -> (array([6, 7, 8, 9]),)

np.where(a % 3 == 0) # 1, 4, 7번 위치가 6, 9, 12로 만족 -> (array([1, 4, 7]),)

0번 부터 시작하는 인덱스 번호에서 몇 번째 위치가 조건을 만족하는지를 반환하고 있습니다.

이번에는 2차원 어레이에 적용하는 예시를 보여드리겠습니다.

a = np.array([[15, 8, 12], [11, 7, 3]])
np.where(a > 10) # (array([0, 0, 1]), array([0, 2, 0]))

다소 결과가 복잡해보이는데, 2차원 이상이면 axis를 기준으로 인덱스 번호를 가져오게 됩니다.

(array([0, 0, 1]), array([0, 2, 0]))를 해석해보면, 앞의 array는 axis = 0, 뒤의 array는 axis = 1에 해당하고,

(0, 0), (0, 2), (1, 0)의 자리가 실제로 각각 15, 12, 11의 값이었으므로, 10보다 크다는 조건과 매칭된다는 점을 알 수 있습니다.

반응형

np.where 값 반환 응용 (조건 True, False 위치 브로드캐스팅)

np.where에서 두번째, 세번째 인자에 True일 때 값, False일 때 값을 써주면 브로드 캐스팅이 가능합니다.

다음의 예제를 살펴보겠습니다.

a = np.array([[15, 8, 12], [11, 7, 3]])
np.where(a > 10, a, 10) # a가 10보다 크면 그대로, 그렇지 않으면 10으로 설정
# array([[15, 10, 12],
#      [11, 10, 10]])

b = np.array([[0, -1, 2], [-2, -1, -3]])
np.where(a > 10, a * 2, b) # a가 10보다 크면 2배, 그렇지 않으면 b의 해당 위치 값으로 설정
# array([[30, -1, 24],
#      [22, -1, -3]])

첫 번째 인자의 조건을 만족하면 해당 위치는 두 번째 인자의 연산으로 브로드 캐스팅되고, 만족하지 않는 경우는 세 번째 인자의 연산으로 브로딩 캐스팅되는 예시를 살펴볼 수 있습니다.

 

이상으로 파이썬 넘파이에서 간단하지만 알면 굉장히 유용한 np.where 메소드의 사용 방법에 대해서 간단히 다루어보았습니다.