Numpy - where() 사용 방법

Python numpy의 where 함수는 배열의 요소가 어떤 조건에 따라 값을 선택하는데 사용됩니다.

예제와 함께 자세히 알아보겠습니다.

1. where()의 Syntax

기본적으로 where()는 condition을 만족하는 배열 요소들의 Index를 리턴합니다.

result = np.where(condition)

다른 사용 방법으로, condition의 true/false 결과에 따라 특정 값으로 변환할 수 있습니다.

• 배열의 요소가 condition에 대해서 true면 x가 리턴되고, false면 y가 리턴
• 각각의 결과들이 모여서 배열로 리턴됨

result = np.where(condition, x, y)

다음으로 설명하는 예제들을 보시면 좀 더 쉽게 이해하실 수 있습니다.

2. 조건을 만족하는 Index 찾기

아래 예제는 배열에서 3보다 큰 요소의 Index를 찾아서 리턴합니다.

• 결과 [3 4]는, 배열의 Index 3과 4에 3보다 큰 값이 위치해 있다는 의미

import numpy as np

arr = np.array([1, 2, 3, 4, 5])
result = np.where(arr > 3)

print(result)
print(result[0])

Output:

(array([3, 4]),)
[3 4]

3. 양수, 음수를 1과 0으로 변경

위와는 다른 사용 방법으로, where()를 사용하여 condition의 true/false 결과에 따라 특정 값으로 변환할 수 있습니다.

아래 예제는 배열의 요소들에 대해서 양수면 1, 음수면 0으로 변환하는 예제입니다.

• 배열의 요소들이 where()의 condition에 대해서 true면 1, false면 0으로 계산
• 결과를 보면 각 요소의 결과가 배열로 저장되어 리턴됨

import numpy as np

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

print(result)

Output:

[0 1 0 1 0]

4. 조건에 따라 값 변경

아래 예제는 요소가 3보다 크면 변경하지 않고, 3 이하 숫자는 0으로 변환하는 예제입니다.

import numpy as np

arr = np.array([1, 2, 3, 4, 5])
result = np.where(arr > 3, arr, 0)

print(result)

Output:

[0 0 0 4 5]

5. 조건에 따라 다른 배열 선택

아래 예제는 condition의 true/false 값에 따라서 다른 배열의 값을 선택하는 예제입니다.

• condition 배열의 어떤 요소 값이 true면 a 배열의 동일 Index 값 선택, 다르면 b 배열의 동일 Index 값 선택

import numpy as np

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

print(result)

Output:

[ 1 20  3 40  5]