LeetCode 703 数据流中的第K大元素 [Swift]
2019-10-29
题目
设计一个找到数据流中第 K 大元素的类（class）。注意是排序后的第 K 大元素，不是第 K 个不同的元素。
你的 KthLargest 类需要一个同时接收整数 k 和整数数组nums 的构造器，它包含数据流中的初始元素。每次调用 KthLargest.add，返回当前数据流中第K大的元素。
示例:
int k = 3;
int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3); // returns 4
kthLargest.add(5); // returns 5
kthLargest.add(10); // returns 5
说明:
你可以假设 nums 的长度≥ k-1 且k ≥ 1。
思路
本题求解思考过程中，我们很容易想到维护一个 size 为 k 的有序数组，每次向数组中添加元素时，如果该元素比数组中的最小值小，跳过；如果比最小值还大，移除最小值，然后将元素加入有序数组。
我们把这个问题转换成了一个排序问题，最容易想到的插入排序，但时间复杂度为 O(n^2)。
我们仔细阅读本题的要求，本题只是需要求解第 K 大元素，另外 K - 1 个元素的排序其实并无要求，只需要比第 K 个元素大即可，我们想到了数据结构：堆（这里我们有必要回归一下堆的相关基础知识），维护一个 size 为 K 的小根堆即可。
现在解题的思路我们已经初见雏形，落地到代码上，我们需要借助利用堆的数据结构实现的 ”优先队列“，优先队列的几个关键方法实现如下：
- 建立：从数组构造优先队列的过程，即为堆的创建过程，从最后一个叶子节点的父节点开始，循环到根节点，依次执行 “上浮” 的操作。
- 入队：首先将入队节点置于最后一个叶子节点的位置，然后对它执行 “上浮” 操作，从而修复因为入队破坏的堆结构。
- 出队：将根节点和最后一个叶子节点互换，将最后一个节点出队，然后对根节点进行 “下沉” 操作，从而修复因为出队破坏的堆结构。
Swift 基础库并没有像 Java，C++ 提供优先队列的数据结构，我们需要基于数组进行封装，具体代码实现如下：
struct Heap <Element> {
var elements : [Element]
let priorityFunction: (Element, Element) -> Bool
init(elements : [Element] = [], priorityFunction : @escaping (Element, Element) -> Bool) {
self.elements = elements
self.priorityFunction = priorityFunction
buildHeap()
}
var isEmpty : Bool {
return elements.isEmpty
}
var count : Int {
return elements.count
}
mutating func buildHeap() {
for index in (0 ..< self.elements.count / 2).reversed() {
siftDown(index)
}
}
func peek() -> Element? {
return elements.first
}
func isRoot(_ index: Int) -> Bool {
return (index == 0)
}
func leftChildIndex(of index: Int) -> Int {
return (2 * index) + 1
}
func rightChildIndex(of index: Int) -> Int {
return (2 * index) + 2
}
func parentIndex(of index: Int) -> Int {
return (index - 1) / 2
}
mutating func equeue(_ element: Element) {
elements.append(element)
siftUP(elementAtIndex: elements.count - 1)
}
mutating func dequeue() -> Element? {
guard !isEmpty else { return nil }
swapElement(at: 0, with: count - 1)
let pop = elements.popLast()
defer {
if !isEmpty {
siftDown(0)
}
}
return pop
}
}
private extension Heap {
func isHigherPriority(firstIndex: Int, secondIndex: Int) -> Bool{
return priorityFunction(elements[firstIndex], elements[secondIndex])
}
func highestPriorityIndex(of parentIndex: Int, and childIndex: Int) -> Int {
guard childIndex < count && isHigherPriority(firstIndex: childIndex, secondIndex: parentIndex)
else { return parentIndex }
return childIndex
}
func highestPriorityIndex(for parent: Int) -> Int {
return highestPriorityIndex(of: highestPriorityIndex(of: parent, and: leftChildIndex(of: parent)), and: rightChildIndex(of: parent))
}
mutating func siftDown(_ elementIndex: Int) {
let highestIndex = highestPriorityIndex(for : elementIndex)
if highestIndex == elementIndex {
return
}
swapElement(at: elementIndex, with: highestIndex)
siftDown(highestIndex)
}
mutating func swapElement(at firstIndex: Int, with secondIndex: Int) {
guard firstIndex != secondIndex
else { return }
elements.swapAt(firstIndex, secondIndex)
}
mutating func siftUP(elementAtIndex: Int) {
let parentIndex = self.parentIndex(of: elementAtIndex)
guard !isRoot(elementAtIndex), isHigherPriority(firstIndex: elementAtIndex, secondIndex: parentIndex) else {
return
}
swapElement(at: elementAtIndex, with: parentIndex)
siftUP(elementAtIndex: parentIndex)
}
}
实现优先队列后，本题可以迎刃而解了，建立堆的时间复杂度为 O(n)，每次入堆的时间复杂度为 O(log N)
class KthLargest {
var heap : Heap<Int>
let k : Int
init(_ k : Int, _ nums : [Int]) {
self.heap = Heap<Int>(elements: nums, priorityFunction : <)
self.k = k
while heap.count > k {
heap.dequeue()
}
}
func add(_ val : Int) -> Int {
if heap.count == k {
if val > heap.peek()! {
heap.dequeue()
heap.equeue(val)
}
} else {
heap.equeue(val)
}
return heap.peek()!
}
}