파이썬 heapq 커스텀 정렬 이용하기

파이썬에서 커스텀 정렬을 이용할 때 리턴값을 True False가 아닌 1, -1로 해주어야 함을 확인하였다.

 

이번에는 여러개의 멤버변수를 가진 클래스를 heapq를 이용해 heap을 만들어줄 때 custom comparator를 어떻게 이용하는지 알아보자.

 

결론부터 말하면

 

 

class 안에 __lt__(self, other) 함수를 재정의 해 주고 return 값은 True False로 해주면 된다.

 

 

이유는 간단하다. heapq의 소스코드를 분석해보면 힙정렬을 할 때 연산기호로 < 를 사용하기 떄문에 __lt__ 재정의 해 주어야 한다

 

 

위에 보면 연산기호가 < 인것을 확인할 수 있다

 

 

import heapq

class node:
    def __init__(self, A, B):
        self.A = A
        self.B = B

    def __lt__(self, other):
        if self.A < other.A:   #오름차순
            return True
        elif self.A == other.A:
            return self.B > other.B  #첫번재 변수가 같으면 두번재 변수로 내림차순
        else:
            return False

    def __str__(self):
        return 'A : {}, B : {}'.format(self.A, self.B)


l = []
heapq.heappush(l, node(1,1))
heapq.heappush(l, node(1,2))

heapq.heappush(l, node(2,1))
heapq.heappush(l, node(2,2))

heapq.heappush(l, node(3,1))
heapq.heappush(l, node(3,2))

heapq.heappush(l, node(4,1))
heapq.heappush(l, node(4,2))

while l:
    print(heapq.heappop(l))

 

첫번째 변수는 오름차순으로, 두번째 변수는 내림차순으로 정렬하도록 만들어 주었다.

 

 

 

결과

A : 1, B : 2
A : 1, B : 1
A : 2, B : 2
A : 2, B : 1
A : 3, B : 2
A : 3, B : 1
A : 4, B : 2
A : 4, B : 1

 

 

 

결론

 

• sort, sorted 함수에 custom comparator를 넣어줄땐 리턴값을 1, -1을 넣어주고
• 클래스에서 custom comparator를 정의 해 주려면 리턴값을 True False로 넣어주자

 

 

 

 

 

근데 사실 이거보다 좋은건 두개의 데이터를 튜플로 묶고, 부호만 바꿔서 정렬해주는게 속도도 빠르고 코드도 짧다