박종훈 알고리즘 블로그

(Leetcode) 23 - Merge k Sorted Lists 풀이

New Year Gift - Curated List of Top 75 LeetCode Questions to Save Your Time

위 링크에 있는 추천 문제들을 시간이 있을때마다 풀어보려고 한다.


/assets/images/2024-02-19-leetcode-23/thumbnail.png

https://leetcode.com/problems/merge-k-sorted-lists/description/

이번 문제는 heap으로 푸는 문제이다. 하지만 문제를 딱 보았을 때 sort로도 풀 수 있을것 같은데? 라는 생각이 들어서 두 가지 방법으로 모두 풀어보았다.

heap으로 풀어보기

python 자체적으로 heap 객체를 쓸 수 있게 제공되는것 같다. 하지만 공부를 위해 직접 구현해서 풀었다. 그러다보니 코드가 엄청 길어졌다.

from typing import Optional
from typing import List
import math


class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next


class Heap:
    def __init__(self):
        self.array = []
        self.size = 0

    def add(self, value: int):
        self.array.append(value)
        self.size += 1

        pointer = self.size
        while pointer > 1:
            temp_pointer = math.floor(pointer/2)
            if self.array[temp_pointer-1] > self.array[pointer-1]:
                temp = self.array[temp_pointer-1]
                self.array[temp_pointer-1] = self.array[pointer-1]
                self.array[pointer-1] = temp

                pointer = temp_pointer
            else:
                break

    def pop(self):
        if self.size == 0:
            return None

        if self.size == 1:
            self.size -= 1
            return self.array.pop()

        result = self.array[0]
        self.array[0] = self.array.pop()
        self.size -= 1

        currentIndex = 0
        lastIndex = self.size - 1

        while True:
            leftIndex = currentIndex * 2 + 1
            rightIndex = currentIndex * 2 + 2

            if lastIndex < leftIndex:
                break

            elif lastIndex == leftIndex:
                if self.array[currentIndex] > self.array[lastIndex]:
                    temp = self.array[currentIndex]
                    self.array[currentIndex] = self.array[lastIndex]
                    self.array[lastIndex] = temp
                currentIndex = leftIndex

            else:
                smallerIndex = leftIndex if self.array[leftIndex] < self.array[rightIndex] else rightIndex
                if self.array[currentIndex] > self.array[smallerIndex]:
                    temp = self.array[currentIndex]
                    self.array[currentIndex] = self.array[smallerIndex]
                    self.array[smallerIndex] = temp
                currentIndex = smallerIndex

        return result


class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        if len(lists) == 0:
            return None

        heap = Heap()
        for list in lists:
            current = list
            while current is not None:
                heap.add(current.val)
                current = current.next

        if heap.size == 0:
            return None

        result = ListNode(val=heap.pop())
        current = result
        while heap.size > 0:
            current.next = ListNode(val=heap.pop())
            current = current.next

        return result

quicksort로 풀어보기

마찬가지로 sort를 가져다 쓰면 편했겠지만 직접 quicksort를 구현하여 풀었다.

import json

from typing import Optional
from typing import List


from typing import Optional
from typing import List


class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next


class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        if len(lists) == 0:
            return None

        concat = []
        for list in lists:
            current = list
            while current is not None:
                concat.append(current.val)
                current = current.next

        if len(concat) == 0:
            return None
        else:
            def quick_sort(arr):
                if len(arr) <= 1:
                    return arr
                pivot = arr[len(arr) // 2]
                left = [x for x in arr if x < pivot]
                middle = [x for x in arr if x == pivot]
                right = [x for x in arr if x > pivot]
                return quick_sort(left) + middle + quick_sort(right)

            sorted = quick_sort(concat)

            head = ListNode(val=sorted.pop(0))
            current = head
            for val in sorted:
                current.next = ListNode(val=val)
                current = current.next
            return head

다른 풀이?

이상하게 푸신분들이 있었다.

class Solution:
    f = open("user.out", 'w')
    print(sys.stdin)
    for s in sys.stdin:
        print(s)
        print('[', ','.join(
            map(str, sorted(int(v) for v in s.rstrip().replace('[', ',').replace(']', ',').split(',') if v))), ']', sep='', file=f)
    exit()

기타

  • 첫 hard 난이도 문제였다.

  • 300만위 안으로 들어왔다.

300만위 안으로 들어왔다

자바로 다시 풀기 (24.07.30)

풀이 1

lists를 순회하면서 가장 작은 수를 가진 노드를 찾고, 그 노드를 mergedList 에 추가한다.

class Solution {
    public ListNode mergeKLists(ListNode[] lists) {
        ListNode mergedList = new ListNode();

        ListNode current = mergedList;

        while (true) {
            int minIndex = -1;
            int currentMin = Integer.MAX_VALUE;

            for (int i = 0; i < lists.length; i++) {
                ListNode node = lists[i];
                if (node != null && node.val < currentMin) {
                    minIndex = i;
                    currentMin = node.val;
                }
            }

            if (minIndex == -1) {
                break;
            }

            current.next = lists[minIndex];
            lists[minIndex] = lists[minIndex].next;

            current = current.next;
        }

        return mergedList.next;
    }
}

TC, SC

문제에서 다음과 같이 정의가 되어있다.

k == lists.length

추가적으로 n을 list 들의 item 수의 총합 이라고 정의하였을 때

시간복잡도는 O(n * k), 공간복잡도는 O(n) 이다.

풀이 2: stream 사용해서 풀기

우선 다 하나의 리스트에 합친 후 정렬한다.

class Solution {
    public ListNode mergeKLists(ListNode[] lists) {
        List<ListNode> mergedListNode = new ArrayList<>();
        for (ListNode listNode : lists) {
            ListNode current = listNode;
            while (current != null) {
                mergedListNode.add(current);
                current = current.next;
            }
        }

        ListNode listNode = new ListNode();
        final ListNode[] current = {listNode};
        mergedListNode.stream().sorted(Comparator.comparingInt(node -> node.val))
                .forEach(node -> {
                    current[0].next = node;
                    current[0] = current[0].next;
                });

        return listNode.next;
    }
}

예상과는 다르게 오히려 이 방식이 더 적은 실행시간으로 완료되었다.

TC, SC

n을 list 들의 item 수의 총합 이라고 정의하였을 때, 시간복잡도는 O(nlogn), 공간복잡도는 O(n) 이다.

categories: 스터디-알고리즘

tags: 파이썬 , 알고리즘 , python , algorithm , Leetcode , heap , quicksort , 자바 , Java