자바 - TreeSet으로 코테 푸는 법

마지막 문제다...

 

TreeSet은 Java에서 SortedSet 인터페이스를 <- 다음 글에서 인터페이스 특집 진행한다.

구현하는 자료구조로, 내부적으로 레드-블랙 트리를 사용하여 요소를 자동으로 정렬합니다. TreeSet의 주요 특징은 데이터가 정렬된 상태로 유지되며, 이를 통해 순서에 의존적인 다양한 문제들을 효율적으로 해결할 수 있다는 점입니다.

TreeSet을 사용해야 하는 문제 유형

1. 정렬된 데이터 유지: 입력 데이터의 순서를 자동으로 정렬해야 하는 경우.
2. 집합 연산: 정렬된 상태로 빠른 집합 연산(교집합, 합집합, 차집합 등)을 수행해야 할 때.
3. 범위 탐색: 특정 범위 내의 데이터를 빠르게 조회하거나 처리해야 할 때.

가장 좋은 문제와 솔루션 예제

문제: 데이터 스트림의 중간값 찾기

문제 설명: 연속적으로 입력되는 데이터 스트림에서 각 단계별로 중간값을 계산하여 반환합니다. 데이터 수가 홀수이면 중간값을, 짝수이면 두 중간값의 평균을 반환합니다.

솔루션 코드

import java.util.Collections;
import java.util.PriorityQueue;
import java.util.TreeSet;

public class MedianFinder {

    private TreeSet<Integer> lower;
    private TreeSet<Integer> higher;

    public MedianFinder() {
        lower = new TreeSet<>(Collections.reverseOrder()); // Max-heap behavior
        higher = new TreeSet<>(); // Min-heap behavior
    }

    public void addNum(int num) {
        if (lower.isEmpty() || num <= lower.first()) {
            lower.add(num);
        } else {
            higher.add(num);
        }

        // Balance the sets to ensure lower has the same number of elements or one more than higher
        if (lower.size() > higher.size() + 1) {
            higher.add(lower.pollFirst());
        } else if (higher.size() > lower.size()) {
            lower.add(higher.pollFirst());
        }
    }

    public double findMedian() {
        if (lower.size() == higher.size()) {
            return (lower.first() + higher.first()) / 2.0;
        } else {
            return lower.first();
        }
    }
}

 

코드 설명

• 두 개의 TreeSet을 사용하여 데이터의 절반을 최대 힙(내림차순 정렬된 TreeSet)으로 관리하고, 나머지 절반을 최소 힙(오름차순 정렬된 TreeSet)으로 관리합니다.
• 새로운 숫자가 추가될 때 적절한 TreeSet에 숫자를 추가하고, 두 TreeSet의 크기를 조정하여 항상 균형을 유지합니다.
• 중간값은 lower의 가장 큰 값(최대 힙의 루트)이거나, 두 TreeSet의 루트 값을 평균내어 계산합니다.

이 방법은 데이터 스트림에서 중간값을 효율적으로 추적할 수 있도록 해주며, TreeSet의 자동 정렬 기능 덕분에 중간값을 빠르게 찾을 수 있습니다. 이런 유형의 문제는 특히 금융, 통계, 데이터 분석 등의 분야에서 매우 유용합니다.