# 41.数据流中的中位数

### 题目

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。

例如，

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

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

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

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

示例 1：

输入：

\["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]

\[\[],\[1],\[2],\[],\[3],\[]]

输出：\[null,null,null,1.50000,null,2.00000]

示例 2：

输入：

\["MedianFinder","addNum","findMedian","addNum","findMedian"]

\[\[],\[2],\[],\[3],\[]]

输出：\[null,null,2.00000,null,2.50000]

限制：

最多会对 `addNum` 、 `findMedian` 进行 `50000` 次调用。

注意：本题与主站 295 题相同：<https://leetcode-cn.com/problems/find-median-from-data-stream/>

> 来源：力扣（LeetCode）
>
> 链接：<https://leetcode-cn.com/problems/shu-ju-liu-zhong-de-zhong-wei-shu-lcof> 著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

### 大小堆

```
class MedianFinder {
    
    var minStack = SwiftPriorityQueue<Int>(hasHigherPriority: { $0 > $1 }, isEqual: { $0 == $1 })
    var maxStack = SwiftPriorityQueue<Int>(hasHigherPriority: { $0 < $1 }, isEqual: { $0 == $1 })
    
    /** initialize your data structure here. */
    init() {

    }
    
    func addNum(_ num: Int) {
        if maxStack.count > minStack.count {
            minStack.enqueue(num)
        }
        else if minStack.count > maxStack.count {
            maxStack.enqueue(num)
        }
        else {
            if let minMax = minStack.array().first, minMax > num {
                minStack.enqueue(num)
            }
            else if let maxMin = maxStack.array().first, maxMin < num {
                maxStack.enqueue(num)
            }
            else {
                maxStack.enqueue(num)
            }
        }
        
        if let minMax = minStack.array().first, let maxMin = maxStack.array().first, minMax > maxMin {
            minStack.dequeue()
            minStack.enqueue(maxMin)
            
            maxStack.dequeue()
            maxStack.enqueue(minMax)
        }
        
//        print("大堆：\(maxStack)")
//        print("小堆：\(minStack)")
//        print("======")
    }

    func findMedian() -> Double {
        if minStack.count == maxStack.count, let a = minStack.peek(), let b = maxStack.peek() {
            return Double((a + b)) / 2.0
        }
        else if minStack.count > maxStack.count {
            return Double(minStack.peek() ?? 0)
        }
        else if maxStack.count > minStack.count {
            return Double(maxStack.peek() ?? 0)
        }
        return 0
    }
}

let obj = MedianFinder()
obj.addNum(-1)
obj.addNum(-2)
obj.addNum(-3)
obj.addNum(-4)
obj.addNum(-5)
obj.addNum(6)

let ret_2: Double = obj.findMedian()
```

![1](https://4193904735-files.gitbook.io/~/files/v0/b/gitbook-x-prod.appspot.com/o/spaces%2F-MI8JgbGh3U6X_oedqkm%2Fuploads%2Fgit-blob-cecf7c456358adfbb1c9273de2c1924e977d985c%2F41-01.png?alt=media)

### Swift - 优先队列

```
import Foundation

public struct SwiftPriorityQueue<Element> {
    private let _hasHigherPriority: (Element, Element) -> Bool
    private let _isEqual: (Element, Element) -> Bool

    private var _elements = [Element]()

    public init(hasHigherPriority: @escaping (Element, Element) -> Bool, isEqual: @escaping (Element, Element) -> Bool) {
        _hasHigherPriority = hasHigherPriority
        _isEqual = isEqual
    }

    public mutating func enqueue(_ element: Element) {
        _elements.append(element)
        bubbleToHigherPriority(_elements.count - 1)
    }

    public func peek() -> Element? {
        return _elements.first
    }
    
    public func array() -> [Element] {
        return _elements
    }

    public var isEmpty: Bool {
        return _elements.count == 0
    }
    
    public var count: Int {
        return _elements.count
    }

    public mutating func dequeue() -> Element? {
        guard let front = peek() else {
            return nil
        }

        removeAt(0)

        return front
    }

    public mutating func remove(_ element: Element) {
        for i in 0 ..< _elements.count {
            if _isEqual(_elements[i], element) {
                removeAt(i)
                return
            }
        }
    }

    private mutating func removeAt(_ index: Int) {
        let removingLast = index == _elements.count - 1
        if !removingLast {
            _elements.swapAt(index, _elements.count - 1)
        }

        _ = _elements.popLast()

        if !removingLast {
            bubbleToHigherPriority(index)
            bubbleToLowerPriority(index)
        }
    }

    private mutating func bubbleToHigherPriority(_ initialUnbalancedIndex: Int) {
        precondition(initialUnbalancedIndex >= 0)
        precondition(initialUnbalancedIndex < _elements.count)

        var unbalancedIndex = initialUnbalancedIndex

        while unbalancedIndex > 0 {
            let parentIndex = (unbalancedIndex - 1) / 2
            guard _hasHigherPriority(_elements[unbalancedIndex], _elements[parentIndex]) else { break }
            _elements.swapAt(unbalancedIndex, parentIndex)
            unbalancedIndex = parentIndex
        }
    }

    private mutating func bubbleToLowerPriority(_ initialUnbalancedIndex: Int) {
        precondition(initialUnbalancedIndex >= 0)
        precondition(initialUnbalancedIndex < _elements.count)

        var unbalancedIndex = initialUnbalancedIndex
        while true {
            let leftChildIndex = unbalancedIndex * 2 + 1
            let rightChildIndex = unbalancedIndex * 2 + 2

            var highestPriorityIndex = unbalancedIndex

            if leftChildIndex < _elements.count && _hasHigherPriority(_elements[leftChildIndex], _elements[highestPriorityIndex]) {
                highestPriorityIndex = leftChildIndex
            }

            if rightChildIndex < _elements.count && _hasHigherPriority(_elements[rightChildIndex], _elements[highestPriorityIndex]) {
                highestPriorityIndex = rightChildIndex
            }

            guard highestPriorityIndex != unbalancedIndex else { break }
            _elements.swapAt(highestPriorityIndex, unbalancedIndex)

            unbalancedIndex = highestPriorityIndex
        }
    }
}

extension SwiftPriorityQueue : CustomDebugStringConvertible {
    public var debugDescription: String {
        return _elements.debugDescription
    }
}
```