题解 | 数据流中的中位数

1、解题思路

直接排序法：每次调用 Insert() 方法时，将数值插入数组并保持数组有序。调用 GetMedian() 时，根据数组长度直接计算中位数。时间复杂度：插入时间复杂度：O(n)（因为插入排序需要移动元素）。获取中位数时间复杂度：O(1)。空间复杂度：O(n)。
堆（优先队列）优化：使用两个堆：一个大根堆（maxHeap）存储较小的半部分数值。一个小根堆（minHeap）存储较大的半部分数值。维护两个堆的大小关系：如果两个堆的大小相等，中位数是两个堆顶的平均值。如果两个堆的大小相差1，中位数是较大堆的堆顶。时间复杂度：插入时间复杂度：O(logn)（堆的插入操作）。获取中位数时间复杂度：O(1)。空间复杂度：O(n)。

2、代码实现

C++

#include <vector>
class Solution {
  public:
    void Insert(int num) {
        nums.push_back(num);
        sort(nums.begin(), nums.end());
    }

    double GetMedian() {
        int n = nums.size();
        if (n % 2 == 1) {
            return nums[n / 2];
        } else {
            return (nums[n / 2 - 1] + nums[n / 2]) / 2.0;
        }
    }

  private:
    vector<int> nums;
};

Java

import java.util.*;


public class Solution {
    private ArrayList<Integer> nums = new ArrayList<>();

    public void Insert(Integer num) {
        nums.add(num);
        Collections.sort(nums);
    }

    public Double GetMedian() {
        int n = nums.size();
        if (n % 2 == 1) {
            return (double) nums.get(n / 2);
        } else {
            return (nums.get(n / 2 - 1) + nums.get(n / 2)) / 2.0;
        }
    }
}

Python

# -*- coding:utf-8 -*-
class Solution:
    def __init__(self):
        self.nums = []

    def Insert(self, num):
        # write code here
        self.nums.append(num)
        self.nums.sort()

    def GetMedian(self):
        # write code here
        n = len(self.nums)
        if n % 2 == 1:
            return self.nums[n // 2]
        else:
            return (self.nums[n // 2 - 1] + self.nums[n // 2]) / 2.0

堆优化进阶版

C++

#include <queue>
#include <vector>
#include <iomanip>
#include <sstream>

class Solution {
  public:
    void Insert(int num) {
        if (maxHeap.empty() || num <= maxHeap.top()) {
            maxHeap.push(num);
        } else {
            minHeap.push(num);
        }

        // 平衡两个堆的大小
        if (maxHeap.size() > minHeap.size() + 1) {
            minHeap.push(maxHeap.top());
            maxHeap.pop();
        } else if (minHeap.size() > maxHeap.size()) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }

    double GetMedian() {
        if (maxHeap.size() == minHeap.size()) {
            return (maxHeap.top() + minHeap.top()) / 2.0;
        } else {
            return maxHeap.top();
        }
    }

  private:
    priority_queue<int> maxHeap;    // 存储较小的半部分 (大根堆)
    priority_queue<int, vector<int>, greater<int>> minHeap; // 存储较大的半部分 (小根堆)
};

Java

import java.util.*;


public class Solution {
    private PriorityQueue<Integer> maxHeap = new PriorityQueue<>
    (Collections.reverseOrder());
    private PriorityQueue<Integer> minHeap = new PriorityQueue<>();

    public void Insert(Integer num) {
        if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
            maxHeap.offer(num);
        } else {
            minHeap.offer(num);
        }

        // 平衡两个堆的大小
        if (maxHeap.size() > minHeap.size() + 1) {
            minHeap.offer(maxHeap.poll());
        } else if (minHeap.size() > maxHeap.size()) {
            maxHeap.offer(minHeap.poll());
        }
    }

    public Double GetMedian() {
        if (maxHeap.size() == minHeap.size()) {
            return (maxHeap.peek() + minHeap.peek()) / 2.0;
        } else {
            return maxHeap.peek().doubleValue();
        }
    }
}

Python

# -*- coding:utf-8 -*-
import heapq


class Solution:
    def __init__(self):
        self.maxHeap = []  # 存储较小的半部分（大根堆）
        self.minHeap = []  # 存储较大的半部分（小根堆）

    def Insert(self, num):
        # write code here
        if not self.maxHeap or num <= -self.maxHeap[0]:
            heapq.heappush(self.maxHeap, -num)
        else:
            heapq.heappush(self.minHeap, num)

        # 平衡两个堆的大小
        if len(self.maxHeap) > len(self.minHeap) + 1:
            heapq.heappush(self.minHeap, -heapq.heappop(self.maxHeap))
        elif len(self.minHeap) > len(self.maxHeap):
            heapq.heappush(self.maxHeap, -heapq.heappop(self.minHeap))

    def GetMedian(self):
        # write code here
        if len(self.maxHeap) == len(self.minHeap):
            return (-self.maxHeap[0] + self.minHeap[0]) / 2.0
        else:
            return -self.maxHeap[0]

3、复杂度分析

直接排序法：每次插入后排序，简单直观。适用于数据量较小的情况。时间复杂度较高（O(n^2) 最坏情况），但能满足题目要求。
堆优化版本：使用两个堆维护数据流的中位数。插入时间复杂度为 O(logn)，获取中位数时间复杂度为 O(1)。适用于数据量较大的情况，效率更高。
堆的平衡：确保两个堆的大小最多相差1。如果两个堆大小相等，中位数是堆顶的平均值。否则，中位数是较大堆的堆顶。
边界条件：数据流为空时如何处理（题目保证数据流至少有一个数值）。数据流长度为1时直接返回该数值。