Ricky.L的随笔
0%

LeetCode.295.数据流的中位数（困难）

发表于 更新于 分类于 Disqus：

题目

LeetCode.295.数据流的中位数（困难）

中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。

例如，

[2,3,4] 的中位数是 3

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

设计一个支持以下两种操作的数据结构：

void addNum(int num) - 从数据流中添加一个整数到数据结构中。
double findMedian() - 返回目前所有元素的中位数。

进阶:

  • 如果数据流中所有整数都在 0 到 100 范围内，你将如何优化你的算法？
  • 如果数据流中 99% 的整数都在 0 到 100 范围内，你将如何优化你的算法？

题解

解法1：排序 + 二分查找

数据结构
列表

如何添加元素？

如果每次添加新元素都要排序，时间复杂度为O(n * logn)。

而列表本身已经是升序，可以用二分查找寻找插入位置，时间复杂度O(log n)，插入后要把插入元素后面的元素全部往后移动一位。

最坏情况下插入到第一个位置，移动所有的数，所以最坏时间复杂度为O(n)。

如何取中位数？

中位数直接取列表的中间的值，时间复杂度为O(1)。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class MedianFinder() {


    /** initialize your data structure here. */
    private val nums = mutableListOf<Int>()

    fun addNum(num: Int) {
        val insertedIndex = binarySearchInsertIndex(num)
        // 每次添加新元素都要把插入元素后面的元素全部往后移动一位，最坏情况下插入到第一个位置，移动所有的数，所以最坏时间复杂度为O(n)
        nums.add(insertedIndex, num)
    }

    // 二分查找插入位置最坏的时间复杂度为O(lgn)
    private fun binarySearchInsertIndex(num: Int): Int {
        val n = nums.size
        var low = 0
        var high = n - 1
        while (low <= high) {
            val mid = low + (high - low) / 2
            if (num <= nums[mid]) {
                if (mid == 0 || num > nums[mid - 1]) {
                    return mid
                } else {
                    high = mid - 1
                }
            } else {
                low = mid + 1
            }
        }
        return n
    }

    // 中位数直接取列表的值，时间复杂度为O(1)
    fun findMedian(): Double {
        val n = nums.size
        return if (n % 2 == 0) (nums[n / 2 - 1] + nums[n / 2]).toDouble() / 2f
        else nums[n / 2].toDouble()
    }
}

解法2：大顶堆 + 小顶堆

数据结构的选择

要保证查询中位数的时间复杂度是O(1)，就要在插入元素的时候就保持元素有序，才能直接查询到中位数。

从数据流中不停的插入新元素并保持有序，用离线排序算法复杂度较高，考虑在线排序算法，用堆排序，API层使用优先队列，添加一个新元素调整堆的时间复杂度最多为O(log n)。

用优先队列怎么获取中位数？

中位数的定义？
假设一共有n个数：

  • n为奇数时，中位数 = 第n / 2 + 1大的数。
  • n为偶数时，中位数 = (第n / 2大的数 + 第n / 2 + 1大的数) / 2

优先队列怎么取元素的？

  • 如果是小顶堆，只能取堆顶最小的元素。
  • 如果是大顶堆，只能取堆顶最大的元素。

我们需要取第n / 2大的数和第n / 2 + 1大的数，但是却只能从堆顶取，那么：

  • 用一个大顶堆存储第1大到第n / 2大的数，堆顶就是第n / 2大的数。
  • 用一个小顶堆存储第n / 2 + 1大到第n大的数，堆顶就是第n / 2 + 1大的数。
  • n为奇数时，大顶堆比小顶堆多存储一个数，堆顶就是中位数。
  • n为偶数时，大顶堆和小顶堆元素个数一样多。

添加新元素时要怎么往优先队列里插入？

初始时，两个堆都为空，允许大顶堆比小顶堆多一个元素，一开始就插入大顶堆。

后续添加新元素，由于是要寻找所有元素中的中位数，所以添加新元素后，中位数可能会变化，两个堆中元素都要调整顺序。

新元素先添加到大顶堆，筛选出最大的数到堆顶，再取出大顶堆的堆顶元素插入小顶堆，这样就保持两个堆顶都是正确的中位数了。

如果后续添加新元素后，元素总数变成奇数，需要在上面的基础上再把小顶堆的顶堆元素取出插入大顶堆中，让大顶堆比小顶堆多一个元素。

复杂度

设元素总数为n。

添加元素时，最坏情况下需要三次插入堆，一次插入堆的时间复杂度是O(log n)，三次就是3 * O(log n)，所以添加元素最坏时间复杂度为O(log n)。

查询中位数时间复杂度O(1)。

空间复杂度O(n)：两个优先队列存储数据流所有的数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import java.util.*

class MedianFinder() {

    /** initialize your data structure here. */
    /**
     * 用大顶堆有序数组的中位数的前半部分
     */
    private val maxHeap = PriorityQueue<Int> { o1, o2 -> o2 - o1 }

    /**
     * 用小顶堆保存有序数组中位数的后半部分
     */
    private val minHeap = PriorityQueue<Int>()

    /**
     * 已保存的数字个数
     */
    private var count = 0

    fun addNum(num: Int) {
        count++
        maxHeap.offer(num)
        minHeap.offer(maxHeap.poll())
        if (count and 1 != 0) maxHeap.offer(minHeap.poll())
    }

    fun findMedian(): Double {
        // 元素总数为奇数，直接取大顶堆堆顶元素
        return if (count and 1 != 0) maxHeap.peek().toDouble()
        else (maxHeap.peek() + minHeap.peek()).toDouble() / 2f
    }
}