np.partition으로 topk 추출하기

np.partion으로 topk를 뽑고 싶을 때

np.partition은 내가 정한 경계를 기준으로 리스팅을 해준다.

 

 

예를 들어서,

A = array([1, 2, 3, 4, 5, 6, 7])

A에서 높은 3가지 숫자를 뽑고 싶을 때는

ind = np.partition(A, -3)

이런식으로 작성하면

ind는

ind = array([1, 2, 3, 4, 5, 6, 7])
ind = array([1, 2, 3, 4, 6, 5, 7])
ind = array([3, 2, 4, 1, 6, 5, 7])

같이 될 수 있다. 

이런식으로 숫자를 오름차순이나 내림차순으로 정하지 않으며, 내가 정한 기준으로 기준을 넘는 것과 못 넘는 것들만 정렬해준다고 생각하면 된다.

 

 

위의 ind 변수의 3가지를 보면 np.partition에서 -3을 줬기 때문에 높은 숫자 3개인 5, 6, 7이 뒤에 있고, 나머지 숫자들은 앞에 있다.

그리고 높은 숫자 3개와 그렇지 않은 숫자 들의 정렬은 어떻게든 될 수 있다는 것을 보여준다.

 

 

이걸 조금 응용하면

topk를 뽑고 싶을 때

topk = np.partition(numpy_array, -topk)[-topk:]

 

이렇게 작성하면 topk를 뽑을 수 있다.