zhangdizhangdi

295. 数据流的中位数 🟥

题目

leetcode 困难

题目

中位数是有序整数列表中的中间值。如果列表的大小是偶数，则没有中间值，中位数是两个中间值的平均值。
例如 arr = [2,3,4] 的中位数是 3 。
例如 arr = [2,3] 的中位数是 (2 + 3) / 2 = 2.5 。

实现 MedianFinder 类:

  • MedianFinder() 初始化 MedianFinder 对象。
  • void addNum(int num) 将数据流中的整数 num 添加到数据结构中。
  • double findMedian() 返回到目前为止所有元素的中位数。与实际答案相差 10-5 以内的答案将被接受。

题解

方法：二分查找

  • 时间复杂度：
    • addNum: 二分查找是 O(logN)，但是数组 splice 插入的最坏时间复杂度是 O(N)。因此单次插入是 O(N)。
    • findMedian: O(1)，直接通过索引取值。
  • 空间复杂度：O(N)
js
class MedianFinder {
  constructor() {
    this.arr = []
  }

  addNum(num) {
    // 如果数组为空，或者 num 比数组最后一个元素还大，直接 push
    if (this.arr.length === 0 || num >= this.arr[this.arr.length - 1]) {
      this.arr.push(num)

      console.log('📌 直接 push', this.arr)
      return
    }

    // 二分查找第一个大于 num 的位置
    let left = 0,
      right = this.arr.length - 1
    while (left <= right) {
      let mid = left + ((right - left) >> 1)

      console.log(
        `left: ${left}: ${this.arr[left]}, right: ${right}: ${this.arr[right]}, mid: ${mid}: ${this.arr[mid]}`,
      )

      if (this.arr[mid] === num) {
        left = mid
        break
      } else if (this.arr[mid] < num) {
        left = mid + 1
      } else {
        right = mid - 1
      }
    }

    // 将 num 插入到保持有序的位置
    this.arr.splice(left, 0, num)

    console.log('📌 插入后', this.arr)
  }

  findMedian() {
    const len = this.arr.length
    const mid = len >> 1 // 相当于 Math.floor(len / 2)
    // 如果是奇数，直接返回中间的元素
    if (len % 2 !== 0) {
      return this.arr[mid]
    }
    // 如果是偶数，返回中间两个元素的平均值
    return (this.arr[mid - 1] + this.arr[mid]) / 2
  }
}

const medianFinder = new MedianFinder()
medianFinder.addNum(1)
medianFinder.addNum(2)
console.log('🌰', medianFinder.findMedian()) // 返回 1.5 ((1 + 2) / 2)
medianFinder.addNum(3)
console.log('🌰', medianFinder.findMedian()) // 返回 2.0
medianFinder.addNum(4)
medianFinder.addNum(1)
console.log('🌰', medianFinder.findMedian()) // 返回 2.0
执行结果
📌 直接 push [ 1 ]
📌 直接 push [ 1, 2 ]
🌰 1.5
📌 直接 push [ 1, 2, 3 ]
🌰 2
📌 直接 push [ 1, 2, 3, 4 ]
left: 0: 1, right: 3: 4, mid: 1: 2
left: 0: 1, right: 0: 1, mid: 0: 1
📌 插入后 [ 1, 1, 2, 3, 4 ]
🌰 2

方法：两个堆

js
// 1. 手写一个基础的堆（支持传入比较函数决定是最大堆还是最小堆）
class Heap {
  constructor(compare) {
    this.data = []
    // 如果 compare(a, b) < 0，说明 a 的优先级高，排在堆顶
    this.compare = compare
  }
  size() {
    return this.data.length
  }
  peek() {
    return this.size() === 0 ? null : this.data[0]
  }
  push(val) {
    this.data.push(val)
    this.bubbleUp(this.size() - 1)
  }
  pop() {
    if (this.size() === 0) return null
    if (this.size() === 1) return this.data.pop()
    const top = this.data[0]
    this.data[0] = this.data.pop()
    this.bubbleDown(0)
    return top
  }
  bubbleUp(index) {
    while (index > 0) {
      const parent = (index - 1) >> 1 // 相当于 Math.floor((index - 1) / 2)
      if (this.compare(this.data[index], this.data[parent]) < 0) {
        this.swap(index, parent)
        index = parent
      } else {
        break
      }
    }
  }
  bubbleDown(index) {
    const lastIndex = this.size() - 1
    while (true) {
      const left = index * 2 + 1
      const right = index * 2 + 2
      let find = index

      if (
        left <= lastIndex &&
        this.compare(this.data[left], this.data[find]) < 0
      ) {
        find = left
      }
      if (
        right <= lastIndex &&
        this.compare(this.data[right], this.data[find]) < 0
      ) {
        find = right
      }
      if (find !== index) {
        this.swap(index, find)
        index = find
      } else {
        break
      }
    }
  }
  swap(i, j) {
    ;[this.data[i], this.data[j]] = [this.data[j], this.data[i]]
  }
}

// 2. 主逻辑：利用双堆实现数据流中位数
class MedianFinder {
  constructor() {
    // 最大堆，保存较小的一半 (b - a < 0 意味着 b > a 时优先，即大顶堆)
    this.maxHeap = new Heap((a, b) => b - a)
    // 最小堆，保存较大的一半 (a - b < 0 意味着 a < b 时优先，即小顶堆)
    this.minHeap = new Heap((a, b) => a - b)
  }

  addNum(num) {
    // 先推入最大堆
    this.maxHeap.push(num)
    // 把最大堆的最大值给最小堆，保证最大堆的元素都小于最小堆
    this.minHeap.push(this.maxHeap.pop())

    // 平衡两个堆的大小（保证 maxHeap 个数 >= minHeap 个数）
    if (this.maxHeap.size() < this.minHeap.size()) {
      this.maxHeap.push(this.minHeap.pop())
    }

    console.log('📌 ⬇️', this.maxHeap.data)
    console.log('📌 ⬆️', this.minHeap.data)
  }

  findMedian() {
    if (this.maxHeap.size() > this.minHeap.size()) {
      return this.maxHeap.peek() // 奇数个时，直接返回最大堆堆顶
    } else {
      return (this.maxHeap.peek() + this.minHeap.peek()) / 2 // 偶数个时，取两个堆顶平均值
    }
  }
}

const medianFinder = new MedianFinder()
medianFinder.addNum(1)
medianFinder.addNum(2)
console.log('🌰', medianFinder.findMedian()) // 返回 1.5 ((1 + 2) / 2)
medianFinder.addNum(3)
console.log('🌰', medianFinder.findMedian()) // 返回 2.0
medianFinder.addNum(4)
medianFinder.addNum(1)
console.log('🌰', medianFinder.findMedian()) // 返回 2.0
执行结果
📌 ⬇️ [ 1 ]
📌 ⬆️ []
📌 ⬇️ [ 1 ]
📌 ⬆️ [ 2 ]
🌰 1.5
📌 ⬇️ [ 2, 1 ]
📌 ⬆️ [ 3 ]
🌰 2
📌 ⬇️ [ 2, 1 ]
📌 ⬆️ [ 3, 4 ]
📌 ⬇️ [ 2, 1, 1 ]
📌 ⬆️ [ 3, 4 ]
🌰 2