[파이썬 알고리즘 인터뷰][정렬] 원점에 K번째로 가까운 점

이 글은 "파이썬 알고리즘 인터뷰 (박상길 지음)"을 읽고 주관적으로 요약한 글입니다. 

---

출처 : https://www.onlybook.co.kr/entry/algorithm-interview

 

문제 정의

원점(0,0)으로부터 유클리드 거리가 가까운 순서 k개의 좌표 출력

 

책에서 구현된 코드

class Solution:
    def kClosest(self, points: List[List[int]], K: int) -> List[List[int]]:
        heap = []
        for (x, y) in points:
            dist = x ** 2 + y ** 2
            heapq.heappush(heap, (dist, x, y))

        result = []
        for _ in range(K):
            (dist, x, y) = heapq.heappop(heap)
            result.append((x, y))
        return result

 

기억해야할 기법

• heapq에 정렬 기준은 튜플의 첫 번째 인자

 

내가 구현한 코드

# 라이브러리 사용
class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        return sorted(points, key=lambda x: ((x[0]**2) + (x[1]**2)))[:k]
        
        
# 머지소트
class Solution:
    def mergesort(self, A):
        def merge(left,right):
            result = []
            l = r = 0
            while l < len(left) or r < len(right):
                if l == len(left):
                    result.append(right[r])
                    r += 1
                elif r == len(right):
                    result.append(left[l])
                    l += 1
                elif (left[l][0] ** 2) + (left[l][1] ** 2) < (right[r][0] ** 2) + (right[r][1] ** 2):
                    result.append(left[l])
                    l += 1
                else:
                    result.append(right[r])
                    r += 1
            return result

        if len(A) < 2:
            return A
        left = self.mergesort(A[:len(A)//2])
        right = self.mergesort(A[len(A)//2:])
        return merge(left, right)

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        return self.mergesort(points)[:k]
        
        
# 퀵소트
class Solution:
    def quicksort(self, A, lo, hi):
        def partition(lo, hi):
            pivot = (A[lo][0] ** 2) + (A[lo][1] ** 2)
            left = i = lo
            right = hi
            while i <= hi:
                val = (A[i][0] ** 2) + (A[i][1] ** 2)
                if left < i and pivot > val:
                    A[left], A[i] = A[i], A[left]
                    left += 1
                elif i < right and pivot < val:
                    A[right], A[i] = A[i], A[right]
                    right -= 1
                else:
                    i += 1

            return left

        if lo < hi:
            pivot = partition(lo, hi)
            self.quicksort(A, lo, pivot - 1)
            self.quicksort(A, pivot + 1, hi)

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        self.quicksort(points,0,len(points)-1)
        return points[:k]

• 그냥 라이브러리인 sort를 사용하면 우선순위 큐보다 약간 빠르다
• 직접 구현한 정렬도 사용해봄
  • 머지소트를 구현해서 돌려봤는데, 속도 차이가 3~4배 차이남
    - heapq/sort : 600~700ms
    - 직접 구현한 머지소트 : 2200ms
  • 잉? 혹시 머지소트가 in-place sort가 아니라서 copy 때문에 그런가?
  • 싶어서 quicksort로 해봤는데 더 느림
    - 2700ms
  • 그렇다면 최적화되지않은 정렬의 구현이 문제일듯
  • heap도 내가 직접 구현한 거 돌리면 느리게 나오려나? 나중에 해봐야겠다