[LeetCode] 295. Find Median from Data Stream（在数据流中寻找中位数）

• Difficulty: Hard
• Related Topics: Heap, Design
• Link: https://leetcode.com/problems/find-median-from-data-stream/

Description

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.
中位数是有序列表的中间值。如果列表的长度是偶数，则没有中间值。此时中位数是中间两个数的平均值。

For example,
例如，

[2,3,4], the median is 3
[2, 3, 4] 的中位数是 3

[2,3], the median is (2 + 3) / 2 = 2.5
[2, 3] 的中位数是 (2 + 3) / 2 = 2.5

Design a data structure that supports the following two operations:
设计一个数据结构，使其能支持以下操作：

• void addNum(int num) - Add a integer number from the data stream to the data structure.
  void addNum(int num) - 往此数据解构里添加一个数。
• double findMedian() - Return the median of all elements so far.
  double findMedian() - 返回此数据结构中目前的中位数。

Example

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3) 
findMedian() -> 2

Follow up

1. If all integer numbers from the stream are between 0 and 100, how would you optimize it?
   如果流里的数据都在 0 到 100 之间，你该如何优化？
2. If 99% of all integer numbers from the stream are between 0 and 100, how would you optimize it?
   如果流里 99% 的数据在 0 到 100 之间，你该如何优化？

Solution

看到题目标签里有个 Heap，想了想，这题的解法应该就是维护堆了。起初的思路是：维护一个大顶堆和一个小顶堆。插入的时候同时往这两个地方插，并维护堆的大小在元素总数的一半。然后代码 WA 了。

class MedianFinder() {
    /** initialize your data structure here. */
    private val minHeap = PriorityQueue<Int>()
    private val maxHeap = PriorityQueue<Int>(compareByDescending { it })
    private var size = 0


    fun addNum(num: Int) {
        size++
        minHeap.offer(num)
        maxHeap.offer(num)
        trimHeap()
    }

    fun findMedian(): Double {
        return if (size % 2 == 0) {
            (minHeap.peek() + maxHeap.peek()) / 2.0
        } else {
            maxHeap.peek().toDouble()
        }
    }

    private fun trimHeap() {
        val targetSize = if (size % 2 == 0) {
            size / 2
        } else {
            size / 2 + 1
        }

        while (minHeap.size != targetSize) {
            minHeap.poll()
        }
        while (maxHeap.size != targetSize) {
            maxHeap.poll()
        }
    }
}

问题出在哪里？翻阅 discussion 后懂了：用两个堆没错，但我维护的方式错了。对于一组数，我们需要做的事将其分为前后两部份（排序后的），中位数就从前半段的最后一个与后半段的第一个中决定。根据堆的性质，不难想到：前半段用大顶堆，后半段用小顶堆。插入元素时，先入大顶堆，再调整两堆，保持大顶堆比小顶堆元素数量相等或多一。代码如下：

class MedianFinder() {
    /** initialize your data structure here. */
    private val small = PriorityQueue<Int>(compareByDescending { it })
    private val large = PriorityQueue<Int>()


    fun addNum(num: Int) {
        large.offer(num)
        small.offer(large.poll())
        if (large.size < small.size) {
            large.offer(small.poll())
        }
    }

    fun findMedian(): Double {
        return if (large.size > small.size) {
            large.peek().toDouble()
        } else {
            (large.peek() + small.peek()) / 2.0
        }
    }
}