[알고스팟] RUNNING MEDIAN 알고리즘 풀이(?)

최근 취업 준비를 하면서 알고리즘 문제를 공부할 필요성을 느껴서 문제를 풀어보고 있다. 몇몇 문제는 내가 상상조차 불가능해서 손도 못대는 경우도 있지만...
일단 노력해보고 있다...허허
너무 간단히 푼 문제말고, 어렵게 느껴졌던 문제나...
새롭게 알게된 경우는 블로그에 풀이?를 올려서 나중에 다시 볼 수 있도록 올려볼 생각이다.

오늘 푼 문제는...

문제

한 수열의 중간값(median)은 이 수열을 정렬했을 때 가운데 오는 값입니다. 
예를 들어 {3,1,5,4,2}를 정렬했을 때 가운데 오는 값은 3이지요.
수열의 길이가 짝수일 때는 가운데 있는 두 값 중 보다 작은 것을 수열의 중간값이라고 정의하도록 합시다.

한 수열의 중간값은 수열에 새로운 수가 추가될 때마다 바뀔 수 있습니다. 
텅 빈 수열에서 시작해서 각 수가 추가될 때마다 중간값을 계산하는 프로그램을 작성하세요. 
예를 들어 3, 1, 5, 4, 2 순서대로 숫자가 추가될 경우 수열의 중간값은 3, 1, 3, 3, 3 순서로 변화합니다.

알고스팟 RUNNING MEDIAN

~ 자세한 입력과 출력 및 문제를 풀어보실 분은 위 주소로 가시면 푸실 수 있습니다. ~

문제를 처음보고 '아 , 정렬해서 가운데 값 보내주면 되겠구나! 간단하네!' 하고

이렇게 풀었었다.....

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int count = Integer.parseInt(br.readLine());

        for (int i = 0; i < count; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine(), " ");
            int length = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            long[] A = new long[length];
            A[0] = 1983;
            long sum =A[0];
            for (int j = 1; j < length; j++) {
                A[j] = (A[j-1] * a + b) % 20090711;
                sum+=findMedian(A, j);
            }
            System.out.println(sum%20090711);
        }
    }

    public static long findMedian(long[] arr, int index) {
        if(index==1) {
            return arr[0];
        }
        long[] subArr = Arrays.copyOfRange(arr, 0, index+1);
        Arrays.sort(subArr);
        if(subArr.length%2 == 0) {
            return subArr[subArr.length/2-1];
        }
        return subArr[subArr.length/2];
    }

}

당연히 결과는 >>>

시간초과였다....허허

이 후에 알고보니 이 문제는 보통은 2개의 Heap으로 구성해서 한 쪽에는 작은 숫자, 한쪽에는 큰 숫자들을 넣고, 이를 중간중간에 Heap의 원소가 한쪽이 더 커지지 않도록 하는 방법을 이용해서 많이 푼다는 것을 알게 되었다.

그래서 이번에는 구글의 도움으로 아래와 같이 Heap을 이용하여 문제를 풀어보았다.

package algo;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class RunningMedian {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int count = Integer.parseInt(br.readLine());

        for (int i = 0; i < count; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine(), " ");
            int length = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            long[] A = new long[length];
            A[0] = 1983;
            long sum = 0;
            for (int j = 1; j < length; j++) {
                A[j] = (A[j - 1] * a + b) % 20090711;

            }
            PriorityQueue<Long> lowers = new PriorityQueue<>( new Comparator<Long>() {
                public int compare(Long a, Long b) {
                    return -1* a.compareTo(b);
                }
            });
            PriorityQueue<Long> highers = new PriorityQueue<>();

            for(int j=0; j<A.length;j++) {
                long number = A[j];
                addNumber(number, lowers, highers);
                rebalance(lowers, highers);
                sum += getMedian(lowers, highers);
            }
            System.out.println(sum % 20090711);
        }

    }

    private static long getMedian(PriorityQueue<Long> lowers, PriorityQueue<Long> highers) {
        PriorityQueue<Long> biggerHeap = lowers.size()>highers.size() ? lowers: highers;
        PriorityQueue<Long> smallerHeap = lowers.size()>highers.size() ? highers: lowers;

        if(biggerHeap.size() == smallerHeap.size()) {
            return lowers.peek();
        } else {
            return biggerHeap.peek();
        }
    }

    private static void addNumber(long number, PriorityQueue<Long> lowers, PriorityQueue<Long> highers) {
        if(lowers.isEmpty() || number<lowers.peek()) {
            lowers.add(number);
        } else {
            highers.add(number);
        }
    }

    private static void rebalance(PriorityQueue<Long> lowers, PriorityQueue<Long> highers) {
        PriorityQueue<Long> biggerHeap = lowers.size()>highers.size() ? lowers: highers;
        PriorityQueue<Long> smallerHeap = lowers.size()>highers.size() ? highers: lowers;

        if(biggerHeap.size()-smallerHeap.size()>=2) {
            smallerHeap.add(biggerHeap.poll());
        }
    }

}

결과는.................

통과!

비록 온전히 내 머리로만 푼 건 아니지만 정답글자를 보는 건 언제나 기분이 좋은 듯하다.

이 코드를 보면

1. lowers와 highers 를 PriorityQueue로 생성한다. 이 때, lowers는 highers와 다르게 작은 값이 맨뒤로 가야하기 때문에 새로운 Comparator를 사용하여 생성하였다.
2. 반복문을 돌며 addNumber , rebalance, getMedian함수를 반복하며 sum 변수에 결과값을 더해준다.
3. sum을 출력한다.

이런 식으로 구조가 나뉘어 있다.

먼저 addNumber함수를 보자.

private static void addNumber(long number, PriorityQueue<Long> lowers, PriorityQueue<Long> highers) {
        if(lowers.isEmpty() || number<lowers.peek()) {
            lowers.add(number);
        } else {
            highers.add(number);
        }
    }

이 함수는 단순히 새로 들어온 숫자가 lowers에 들어갈지 highers에 들어갈 지를 결정해준다. 각 Heap의 원소의 갯수는 전혀 상관하지 않는다. ( 다음에 나올 rebalance함수에서 처리해 줄 것이다. )

이제 rebalance함수를 보면

private static void rebalance(PriorityQueue<Long> lowers, PriorityQueue<Long> highers) {
        PriorityQueue<Long> biggerHeap = lowers.size()>highers.size() ? lowers: highers;
        PriorityQueue<Long> smallerHeap = lowers.size()>highers.size() ? highers: lowers;

        if(biggerHeap.size()-smallerHeap.size()>=2) {
            smallerHeap.add(biggerHeap.poll());
        }
    }

여기서는 두 Heap중 더 큰 Heap과 작은 Heap을 찾아내고, 각각의 Heap의 원소의 갯수가 2개이상 차이가 나면( 2개이상 원소갯수가 차이가 나면 중간값을 찾을수가 없으므로 )
더 큰 Heap에서 Head에 있는 값을 더 작은 Heap으로 이동시키다.

이제 각각의 Heap에 각각 작은 값들과 큰 값들이 담겨있으므로 중간값을 구하기 매우 수월해졌다.

마지막으로 getMedian함수를 보면

private static void rebalance(PriorityQueue<Long> lowers, PriorityQueue<Long> highers) {
        PriorityQueue<Long> biggerHeap = lowers.size()>highers.size() ? lowers: highers;
        PriorityQueue<Long> smallerHeap = lowers.size()>highers.size() ? highers: lowers;

        if(biggerHeap.size()-smallerHeap.size()>=2) {
            smallerHeap.add(biggerHeap.poll());
        }
    }

여기서는 두 Heap의 사이즈가 같은 경우에는 문제에서 둘중 더 작은 값을 제출하라고 하였기에 lowers에서 값을 가져와 return해주었고, 그렇지 않은 경우에는 원소가 하나 더 많은 Heap에서 peek() 해서 return을 해주었다.

이렇게 하여 문제를 풀었다.

참고로 위 코드는 유튜브 Data Structures: Solve 'Find the Running Median' Using Heaps 를 보면 더 이해가 쉬울 것이다. ~ 코드는 사실 거의 같다... ~

---

아직 모르는게 많아 게시글에 잘못된 정보가 있을 수 있습니다. 혹시 잘못된 정보가 있다면, 댓글 혹은 메일로 알려주시면 최대한 빨리 수정하겠습니다!