파이썬 heapq 모듈을 이용해서 priority queue 를 사용해 봅시다.

 안녕하세요. 5일 만에 돌아온 코딩개 입니다. 제가 낸 코딩 테스트 문제 중에 heap 자료구조를 이용하는 문제가 있었습니다. 1회에 유독 많이 냈던 걸로 기억하는데요. 파이썬에서는 heapq 모듈과 tuple을 이용하면 날먹 수준으로 쉽게 구현하실 수 있어요.

 

---

 예제 코드를 볼게요.

 

 먼저 2번째 줄은 prioroity queue 역할을 할 배열을 선언한 거에요. 다음에 heappush를 하는데요. pq에 tuple (1, 2)를 넣습니다. 4번째 줄에는 (-1, 3)을 넣어요. 다음에, 5번째 줄에서, pq에서 맨 위에 있는 원소를 pop 하면서 맨 위에 있었던 원소를 가져오게 됩니다. 그리고 다시 (-1, 6)을 heap에 넣고, 7번째 줄에서 pq[0]에 무엇이 들어있는지만 봅니다. heappop하고 pq[0]하고 뭐가 다를까요?

 

 전자는 맨 위에 있는 원소를 pop 하면서 가져오지만, 후자는 그냥 맨 위에 있는 원소를 가져오기만 합니다. 제거를 하지 않아요.

 

 

 실행 결과는 위와 같아요. 이게 어떤 원리로 저렇게 들어갔는지를 물어보신다면, 코드 내부를 살짝 보시는 수 밖에 없겠네요.

 

 

 item이 들어오고, append가 되는데요. 예를 들어서, pq에 (1, 2)만 있었는데, (-1, 3)이 들어왔다고 해 봅시다.

 

 

 그러면, 일단 요래 들어옵니다. 다음에 133번째 줄이 수행되는데요. startpos가 0이고 pos가 1인 상황입니다.

 

 

 여기서 newitem은 맨 뒤에 있는 (-1, 2)를 의미해요. while loop를 따라가 봅시다. 먼저 pos인 1이 startpos인 0보다 크므로, while 루프를 타게 되는데요.

 

 

 parent는 노란색을, pos는 군청색을 의미합니다. 212번째 줄의 newitem은 (-1, 3), parent는 (1, 2)인데요. tuple에서 (-1, 3) < (1, 2)이므로, 212번째 줄의 if문에 걸려 들어갑니다.

 

 

 그 다음에 while loop를 돌면 pos가 0이 되는데요. startpos가 0이니까 while 조건에 걸리지 않아요. 따라서 pq의 0번째 위치에 (-1, 3)이 들어갑니다.

 

 

 그래서, 처음에 (1, 2)와 (-1, 3)을 heapq에 넣고 빼면, (-1, 3)이 먼저 나오게 됩니다.

 

 

---

 그러면 이것을 이용해서, 이 문제를 어떻게 풀까요? 우선 순위가 큰 프로세스, 우선 순위가 같다면 id가 작은 게 우선입니다. tuple에 우선 순위, id, 남은 실행 시간으로 넣어도 큰 문제는 없는데요. 왜냐하면 프로세스의 id는 모두 다르기 때문입니다. 우선 순위와 id가 모두 같은 서로 다른 프로세스가 있다면 문제가 생길 수 있지만, 그런 경우는 없다고 주어졌어요.

 

 

 3개의 데이터를 넣어봅시다. 튜플에 (a, b, c)가 들어가는데요. 우선 순위가 a이고, id가 b이고 실행 시간이 c인 프로세스를 의미합니다. 3개를 모두 넣은 다음에 heappop을 시켜 봅시다.

 

 

 그러면 (1, 2, 5)가 나오는데요. 이는 튜플에서 (1, 2, 5) < (1, 3, 2) < (2, 3, 2)이기 때문입니다. id 순으로는 나온 거 같은데, 우선 순위가 낮은 게 먼저 나왔네요? 우선 순위가 2였다, 1이였다는 정보는 보존하되, 우선 순위가 2인 프로세스가 먼저 나오게끔 하는 방법이 없을까요? a < b이면 -b < -a 라는 점을 이용하면 꽤 간단하게 해결할 수 있습니다. 즉, 프로세스의 우선순위에 -1을 곱한 값을 넣어버리면 됩니다.

 

 

 이러면 어떤가요? 우선 순위가 2이면서 id가 3이고, 실행 시간이 2인 프로세스를 (2, 3, 2)로 넣는 대신에 (-2, 3, 2)로 넣습니다.

 

 

 그러면 (-2, 3, 2) < (-1, 3, 2) < (-1, 2, 5)니까 우선 순위가 2인 프로세스가 먼저 나와버리게 됩니다. 물론, 프로세스 모델 class를 만들고 lt를 재정의 하는 방법도 있습니다만, 이보다는 tuple을 이용하는 게 훨씬 간결하고 깔끔합니다.