Python numpy : where() — 조건에 따라 두 값 중 하나를 선택하는 메서드

NumPy where() — 조건에 따라 두 값 중 하나를 선택하는 메서드

조건이 True인 위치에는 x값, False인 위치에는 y값을 선택해서 새 array를 만듭니다.

기본 문법

# 형태 1 — 조건에 따라 값 선택 (가장 많이 쓰는 형태)
numpy.where(condition, x, y)

# 형태 2 — 조건이 True인 인덱스 반환
numpy.where(condition)

형태 1 — 조건에 따라 값 선택

import numpy as np

arr = np.array([1, -2, 3, -4, 5])

np.where(arr > 0, arr, 0)
# array([1, 0, 3, 0, 5])
#           ↑     ↑
#        -2, -4 → 0으로 대체

조건이 True인 위치 → arr 값 그대로, False인 위치 → 0

x, y 자리에 값 대신 array도 가능

a = np.array([1, 2, 3, 4, 5])
b = np.array([10, 20, 30, 40, 50])

np.where(a > 3, a, b)
# array([10, 20, 30,  4,  5])
#         ↑   ↑   ↑   ↑   ↑
# 1≤3→b  2≤3→b  3≤3→b  4>3→a  5>3→a

형태 2 — True인 인덱스 반환

인수를 하나만 넘기면 조건을 만족하는 위치의 인덱스를 반환합니다.

arr = np.array([10, 50, 30, 80, 20])

np.where(arr > 40)
# (array([1, 3]),)  ← 인덱스 1, 3이 조건 만족 (50, 80)

# 인덱스로 값 추출할 때
arr[np.where(arr > 40)]
# array([50, 80])

pandas.where vs numpy.where — 헷갈리기 쉬운 차이

이름은 같지만 동작 방향이 반대입니다.

s = pd.Series([1, -2, 3, -4, 5])

# pandas.where → True면 원래 값 유지, False면 대체
s.where(s > 0, 0)
# [1, 0, 3, 0, 5]   ← 양수는 유지, 음수는 0

# numpy.where → True면 x, False면 y
np.where(s > 0, s, 0)
# [1, 0, 3, 0, 5]   ← 결과는 같지만 사고방식이 다름

	pandas.where(cond, other)	numpy.where(cond, x, y)
True일 때	원래 값 유지	x 선택
False일 때	other로 대체	y 선택