Python으로 데이터 스트림에서 K번째로 큰 요소 찾기

문제 개요

데이터 스트림에서 k번째로 큰 요소를 실시간으로 찾는 클래스를 설계해야 한다고 가정해 보겠습니다. 여기서 말하는 k번째로 큰 요소란 정렬된 순서 기준의 k번째 요소를 의미하며, 중복을 제거한 k번째 고유(distinct) 값이 아니라는 점에 유의해야 합니다.

KthLargest 클래스는 다음과 같은 구조로 동작합니다.

• 생성자(__init__)는 정수 k와 스트림의 초기 데이터가 담긴 배열 nums를 인자로 받습니다.
• add(val) 메서드가 호출될 때마다 새로운 값이 스트림에 추가되며, 현재 스트림에서 k번째로 큰 요소를 반환합니다.

예를 들어, k = 3이고 초기 배열이 [4, 5, 8, 2]일 때, 순서대로 add(3), add(5), add(10), add(9), add(4)를 호출하면 출력은 각각 4, 5, 5, 8, 8이 됩니다.

풀이 접근 방식

가장 직관적인 방법은 새로운 값이 들어올 때마다 배열 전체를 정렬한 뒤, 뒤에서 k번째 위치의 요소를 꺼내는 것입니다. 구현 단계는 다음과 같습니다.

1. 초기화(__init__): k와 nums를 멤버 변수로 저장합니다.
2. add() 함수 정의: 새 값 val을 인자로 받아 처리합니다.
   • val을 배열의 끝에 추가(append)합니다.
   • 배열을 오름차순으로 정렬(sort)합니다.
   • array[len(array) - k], 즉 뒤에서 k번째 요소를 반환합니다.

구현 예제

class KthLargest:
    def __init__(self, k, nums):
        self.array = nums
        self.k = k

    def add(self, val):
        self.array.append(val)
        self.array.sort()
        return self.array[len(self.array) - self.k]

ob = KthLargest(3, [4, 5, 8, 2])
print(ob.add(3))  # 4
print(ob.add(5))  # 5
print(ob.add(10)) # 5
print(ob.add(9))  # 8
print(ob.add(4))  # 8

입력 및 실행 결과

4
5
5
8
8

동작 원리 상세 분석

초기 배열 [4, 5, 8, 2]에서 k = 3인 경우를 단계별로 살펴보겠습니다.

• add(3): 배열은 [2, 3, 4, 5, 8]로 정렬되고, 뒤에서 3번째 값은 4입니다.
• add(5): 배열은 [2, 3, 4, 5, 5, 8]이 되고, 뒤에서 3번째 값은 5입니다.
• add(10): 배열은 [2, 3, 4, 5, 5, 8, 10]이 되고, 뒤에서 3번째 값은 여전히 5입니다.
• add(9): 배열은 [2, 3, 4, 5, 5, 8, 9, 10]이 되고, 뒤에서 3번째 값은 8입니다.
• add(4): 배열은 [2, 3, 4, 4, 5, 5, 8, 9, 10]이 되고, 뒤에서 3번째 값은 8입니다.

중복 값도 하나의 요소로 취급하기 때문에 결과가 위와 같이 나오는 것을 확인할 수 있습니다.

효율성 개선: 힙(Heap) 활용

위 방법은 add()가 호출될 때마다 정렬을 수행하므로 시간 복잡도가 O(n log n)입니다. 스트림의 크기가 커질수록 비효율적일 수 있습니다. 이 경우 최소 힙(min-heap)을 사용하면 훨씬 효율적으로 문제를 해결할 수 있습니다. 힙의 크기를 항상 k로 유지하면, 힙의 루트(root)가 곧 k번째로 큰 요소가 됩니다.

import heapq

class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.heap = []
        for num in nums:
            self.add(num)

    def add(self, val):
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]

이 방식에서 각 add() 호출은 O(log n) 시간 안에 처리되므로, 대규모 스트림 환경에서도 안정적인 성능을 기대할 수 있습니다.

마무리

정렬 기반 풀이는 구현이 간단하고 직관적이라는 장점이 있으며, 힙 기반 풀이는 반복적인 삽입 연산이 많은 실시간 스트림 환경에 더 적합합니다. 문제의 입력 규모와 호출 빈도를 고려하여 적절한 방식을 선택하는 것이 좋습니다.