np.where - 배열 원소에 대한 조건부 검색을 한번에 하자

 

np.where는 배열 원소에 대한 조건부 검색이다.
리턴값은 조건에 맞는 배열의 인덱스 또는 값이다.
리턴값이 튜플인지 ndarray인지 주의해야한다.

 

1.  어떤 배열에 대해 조건에 맞는 인덱스 뽑아내기

np.where(array에 대한 조건)
: 조건이 True에 해당하는 인덱스로 구성된 ndarray를 튜플로 감싸서 반환

 

먼저 1차원 배열에 대한 조건에 맞는 인덱스를 뽑아내는 예제를 만들었다.

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

c = np.where(a > 3)         # 배열 a에 대해 조건이 참인 인덱스(3, 4)를 반환
         
print("반환값 : ", c)
print("반환값의 타입 : ", type(c))
print("반환된 튜플의 첫번째 원소의 타입 :", type(c[0]))
print("반환된 튜플의 첫번째 원소의 차원 :", c[0].shape)

 

output :

반환값 : (array([3, 4]),)
반환값의 타입 : <class 'tuple'>
반환된 튜플의 첫번째 원소의 타입 : <class 'numpy.ndarray'>
반환된 튜플의 첫번째 원소의 차원 : (2,)

아래에 출력된 결과를 보면 리턴값은 튜플이다.

따라서 1차원 array에 대해 조건에 맞는 인덱스를 뽑아낼때는 c[0]으로 튜플의 첫번째 값을 뽑아내야하며,

튜플의 첫번째 값인 c[0]는 조건에 맞는 인덱스를 모아놓은 (2, )의 ndarray임을 알 수 있다.

 

 

다음으로 2차원 배열에 대해 조건에 맞는 인덱스를 뽑아내는 예제를 보자

a = np.array([[1,2,3,4,5],       # (2, 5) array
              [5,4,3,2,1]])       

# (1, )차원의 배열 b를 (2, 5) 차원으로 broadcasting 후에 원소별 조건 검토
b = np.array(3)                  
c = np.where(a > b)            


# 배열 a에 대해 조건이 참인 인덱스를 차원별로 각각 리스트로 반환
# 배열 a에 대해 조건이 참인 첫번째 차원의 인덱스(0, 0, 1, 1) 반환
# 배열 a에 대해 조건이 참인 두번째 차원의 인덱스(3, 4, 0, 1) 반환


print("반환값 : ", c)
print("반환값의 타입 : ", type(c))
print("반환된 튜플의 첫번째 원소의 타입 :", type(c[0]))
print("반환된 튜플의 첫번째 원소의 차원 :", c[0].shape)
print("반환된 튜플의 두번째 원소의 타입 :", type(c[1]))
print("반환된 튜플의 두번째 원소의 차원 :", c[1].shape)

반환값 : (array([0, 0, 1, 1]), array([3, 4, 0, 1]))
반환값의 타입 : <class 'tuple'>
반환된 튜플의 첫번째 원소의 타입 : <class 'numpy.ndarray'>
반환된 튜플의 첫번째 원소의 차원 : (4,)
반환된 튜플의 두번째 원소의 타입 : <class 'numpy.ndarray'> 반환된 튜플의 두번째 원소의 차원 : (4,)

아래에 출력된 결과를 보면 리턴값은 역시 튜플이며 두개의 ndarray로 구성되어 있다.

튜플의 첫번째 ndarray는 a의 첫번째 차원인 (2, )에 대해 조건에 맞는 인덱스를 모아놓은 배열이며

튜플의 두번째 ndarray는 a의 두번째 차원인 (5, )에 대해 조건에 맞는 인덱스를 모아놓은 배열이다.

 

2.  어떤 배열에 대해 조건에 맞는 값 뽑아내기

np.where를 이용하면 배열 a에 대해 원소별로 특정 조건을 만족하는 값을 뽑아낼 수도 있다.

np.where(array에 대한 조건, a, b)
: array에 대한 조건의 각 원소별로 True인 경우 a의 해당 원소를 반환, False인 경우 b의 해당 원소를 반환

 

a = np.array([[1,2,3,4,5],      # (2, 5) array
              [5,4,3,2,1]])   
              
# (2, 5) 차원으로 broadcasting 후에 원소별 조건 검토    
b = np.array(3)                 

c = np.where(a > b, a, -1)      

# 배열 a에 대해 원소별로 b보다 큰 경우 배열 a의 해당 원소를 출력
# 배열 a에 대해 원소별로 b보다 작을 경우     
                                

print("반환값 : ", c)
print("반환값의 타입 : ", type(c))
print("반환값의 차원 : ", c.shape)

반환값 : [[-1 -1 -1 4 5] [ 5 4 -1 -1 -1]]
반환값의 타입 : <class 'numpy.ndarray'>
반환값의 차원 : (2, 5)

 

1.에서 살펴봤던 것과 다르게 어떤 조건에 맞는 값을 뽑아낼때의 리턴값은 차원이 (2, 5)인 ndarray이다.