Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python NumPy에서 특정 범위 내의 배열 요소 찾는 방법

NumPy 라이브러리를 사용해 데이터를 처리하다 보면, 특정 숫자 범위에 해당하는 요소만 골라내야 하는 경우가 자주 발생합니다. 이럴 때 NumPy에서 기본으로 제공하는 함수들을 활용하면 간단하게 조건에 맞는 요소를 필터링할 수 있습니다.

이 글에서는 대표적인 두 가지 방법을 소개합니다. 하나는 np.wherenp.logical_and를 함께 사용하는 방법이고, 다른 하나는 *(곱셈) 연산자를 이용해 불리언 마스크를 만드는 방법입니다.

방법 1: np.logical_and와 np.where 사용하기

이 방식은 NumPy 배열에 logical_and 함수를 적용하여 두 개의 조건(범위의 최솟값 이상, 최댓값 이하)을 동시에 검사합니다. np.where 절과 결합하면, 주어진 범위 조건을 만족하는 요소들의 인덱스 위치를 배열 형태로 얻을 수 있습니다.

import numpy as np

A = np.array([5, 9, 11, 4, 31, 27, 8])

# 초기 배열 출력
print("Given Array : ", A)

# 범위: 6 이상 15 이하
res = np.where(np.logical_and(A >= 6, A <= 15))

# 결과 출력
print("Array with condition : ", res)

실행 결과

위 코드를 실행하면 다음과 같은 결과를 확인할 수 있습니다.

Given Array : [ 5 9 11 4 31 27 8]
Array with condition : (array([1, 2, 6], dtype=int32),)

출력 결과를 보면 인덱스 1, 2, 6 위치에 있는 값들, 즉 9, 11, 8이 조건(6~15 사이)을 만족한다는 것을 알 수 있습니다.

방법 2: * 연산자로 불리언 마스크 만들기

두 번째 방법은 * 연산자를 활용하는 것입니다. 각 비교 연산은 True/False로 구성된 불리언 배열을 반환하는데, 두 배열을 곱하면 논리 AND 연산과 동일한 효과를 냅니다. 이 방식의 장점은 인덱스가 아닌 실제 값 자체를 결과로 얻을 수 있다는 점입니다.

import numpy as np

A = np.array([5, 9, 11, 4, 31, 27, 8])

# 초기 배열 출력
print("Given Array : ", A)

# 범위: 6 이상 15 이하
res = A[(A >= 6) * (A <= 15)]

# 결과 출력
print("Array with condition : ", res)

실행 결과

위 코드를 실행하면 다음과 같은 결과가 출력됩니다.

Given Array : [ 5 9 11 4 31 27 8]
Array with condition : [ 9 11 8]

정리

두 방법 모두 NumPy 배열에서 특정 범위의 값을 효율적으로 추출할 수 있습니다.

  • np.where + np.logical_and: 조건을 만족하는 요소의 인덱스가 필요할 때 유용합니다.
  • * 연산자: 조건을 만족하는 요소의 실제 값을 바로 얻고 싶을 때 코드가 더 간결해집니다.

참고로, * 연산자 대신 & 연산자를 사용해도 동일한 결과를 얻을 수 있으며, 이 경우 각 조건을 반드시 괄호로 묶어주어야 합니다. 상황에 맞게 적절한 방법을 선택해 사용하시기 바랍니다.