[PS] Leet | 295. 데이터 스트림에서 중간값 찾기

문제

The median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value, and the median is the mean of the two middle values.

• For example, for arr = [2,3,4], the median is 3.
• For example, for arr = [2,3], the median is (2 + 3) / 2 = 2.5.

Implement the MedianFinder class:

• MedianFinder() initializes the MedianFinder object.
• void addNum(int num) adds the integer num from the data stream to the data structure.
• double findMedian() returns the median of all elements so far. Answers within 10-5 of the actual answer will be accepted.

예시

Input

["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]

Output

[null, null, null, 1.5, null, 2.0]

Explanation

MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1, 2]
medianFinder.findMedian(); // return 1.5 (i.e., (1 + 2) / 2)
medianFinder.addNum(3);    // arr[1, 2, 3]
medianFinder.findMedian(); // return 2.0

조건

• -105 <= num <= 105
• There will be at least one element in the data structure before calling findMedian.
• At most 5 * 104 calls will be made to addNum and findMedian.

답

가장 먼저 생각나는 방법이 이진 탐색 후 리스트에 항목 추가입니다.

 

하지만 이 경우 탐색에는 O(logn)이 걸리지만 리스트에 항목 추가시 O(n)의 시간이 걸리게 됩니다.

 

리스트 대신 heap 구조에 추가하면 O(logn)의 시간을 유지할 수 있습니다.

 

heap 클래스

오름차순으로 요소를 배열할 heap과 내림차순으로 요소를 배열할 heap 두 가지를 준비합니다.

 

arrUp은 중앙을 기준으로 큰 수들이 있을 우측 배열을 나타내고,

 

arrDown은 중앙을 기준으로 작은 수들이 있을 좌측 배열을 나타냅니다.

 

push는 각 heap의 최하단에 요소를 삽입한 후, 대소관계가 맞을 때 까지 parent를 타고 올라갑니다.

 

pop은 각 heapdml 최상단 요소를 뽑아낸 후, 최하단 요소를  최상단에 두어 대소관계가 맞을 때 까지 child를 타고 내려갑니다.

var MedianFinder = function(val) {
    this.arrUp = [];    // 오름차순 정렬된 heap
    this.arrDown = [];  // 내림차순 정렬된 heap

    // push : 각 heap의 최하단에 요소를 삽입후
    // 대소관계가 맞을 때 까지 parent를 타고 올라감
    this.push = (isUp, num) => {
        const arr = isUp? this.arrUp : this.arrDown;
        let i = arr.length;
        arr.push(num);

        while (i > 0){
            const parentIdx = Math.floor((i - 1) / 2);
            if (isUp && (arr[parentIdx] > num) || !isUp && (arr[parentIdx] < num)) {
                [arr[parentIdx], arr[i]] = [arr[i], arr[parentIdx]];
                i = parentIdx;
            }
            else break;
        }
    }
    
    // pop : 최상단 요소를 뽑아낸 다음 최하단의 요소를 최상단에 둔 후
    // 대소관계가 맞을 때 까지 child를 타고 내려감
    this.pop = (isUp) => {
        const arr = isUp? this.arrUp : this.arrDown;
        if (arr.length === 0) return null

        let i = 0;
        const ret = arr[0];
        const last = arr.pop();
        if (arr.length > 0) arr[0] = last;

        while (true) {
            const childIdx1 = i * 2 + 1;
            const childIdx2 = i * 2 + 2;
            let childIdx = i
            if (childIdx1 < arr.length && (isUp && arr[childIdx1] < arr[childIdx] || !isUp && arr[childIdx1] > arr[childIdx])){
                childIdx = childIdx1;
            }
            if (childIdx2 < arr.length && (isUp && arr[childIdx2] < arr[childIdx] || !isUp && arr[childIdx2] > arr[childIdx])){
                childIdx = childIdx2;
            }
            if (childIdx === i) break;
            [arr[childIdx], arr[i]] = [arr[i], arr[childIdx]];
            i = childIdx;
        }

        return ret
    }
};

addNum ( )

arrUp에 요소를 넣어 한 번 정렬 => 우측 배열에서의 최소값이 나옵니다.

 

해당 최소값을 arrDown에 넣어 한 번 정렬 => 좌측 배열에서의 최대값이 나옵니다.

 

이렇게 하면 arrUp의 최소값과 arrDown의 최대값은 각각 중앙에 있는 값이 됩니다.

/** 
 * @param {number} num
 * @return {void}
 */
MedianFinder.prototype.addNum = function(num) {
    this.push(true, num);
    this.push(false, this.pop(true))

    if (this.arrUp.length < this.arrDown.length) this.push(true, this.pop(false));
    //console.log(this.arrDown, this.arrUp);
};

findMedian ( )

arrUp의 최소값과 arrDown의 최대값 모두 각 heap의 첫 번째 요소를 뽑으면 구할 수 있습니다.

 

(arrUp은 오름차순, arrDown은 내림차순이므로)

 

arrUp이 arrDown보다 한 칸 더 길다면 전체 요소 개수가 홀수이므로 arrUp의 최소값이 중간값이 됩니다.

 

arrUp과 arrDown의 길이가 같아면 전체 요소 개수가 짝수이므로 두 말단 값의 평균이 중간값이 됩니다.

/**
 * @return {number}
 */
MedianFinder.prototype.findMedian = function() {
    const upPeek = this.arrUp[0];
    const downPeek = this.arrDown[0];
    return (this.arrUp.length > this.arrDown.length)? upPeek : (upPeek + downPeek) / 2;
};