【二叉树】数据流中的第 K 大元素

0x00 题目

设计一个找到数据流中第 k 大元素的类（class）
注意是排序后的第 k 大元素，不是第 k 个不同的元素

请实现 KthLargest 类：

KthLargest(int k, int[] nums)
使用整数 k 和整数流 nums 初始化对象

int add(int val)
将 val 插入数据流 nums 后，返回当前数据流中第 k 大的元素

---

0x01 思路

方法一：
先对初始化的数组进行排序
nums.sorted()
时间复杂度：O(NlogN) (快排)

插入新数据时，使用二分搜索法，找到对应的位置，然后插入
nums.insert(val, at: index)

最后返回指定位置的值
nums[nums.count - k]

---

方法二：
使用小顶堆进行数据管理

小顶堆：
每个结点的值都 小于或等于 其左右子节点的值

堆其实就是利用完全二叉树的结构来维护的一维数组

用公式来描述小顶堆就是：
arr[i] <= arr[2i+1] && arr[i] <= arr[2i+2]

构建一个包含 k 个元素的小顶堆
堆顶就是第 k 大的元素

---

0x02 解法

语言：Swift

解法：

class KthLargest {
	// 存储 k 个元素
    private var arr: [Int] = []
    private var k: Int = 0
    
    init(_ k: Int, _ nums: [Int]) {
        self.k = k
        for num in nums {
            _ = add(num)
        }
    }
    
    func add(_ val: Int) -> Int {
    	// 不够 k 个元素，继续添加
        if arr.count < k {
            arr.append(val)
            // 添加元素是追加到末尾，然后进行 "上浮" 操作，把小元素 "浮" 到上面去
            minHeapFixup(&arr)
        }else{
        	// 已经有 k 个元素了，进行判断
            if val > arr[0]  {
            	// 大于就替换，堆顶元素
                arr[0] = val
            }
            // 替换后，然后进行 "下沉" 操作，把大元素 "沉" 到下面去
            minHeapFixdown(&arr, 0, k)
        }
        return arr[0]
    }
    
    //  "上浮" 操作
    private func minHeapFixup(_ arr: inout [Int]) {
        var j = arr.count - 1
        // 索引 (j-1)/2 的节点，是索引 j 节点，的父节点
        // 小的元素 "上浮"，然后再比较 父节点 与 父节点的父节点
        while (j-1)/2 >= 0 && arr[j] < arr[(j-1)/2] {
            arr.swapAt(j, (j-1)/2)
            j = (j-1)/2
        }
    }
    
    // "下沉" 操作
    private func minHeapFixdown(_ arr: inout [Int], _ i: Int, _ n: Int) {
        var p = i // 父节点索引
        let count = arr.count
        
        while true {
            var min = p

            // 根节点与左节点比较，如果大于左节点，记录索引
            if p*2+1 < count && arr[p] > arr[p*2+1] {
                min = p*2 + 1
            }
            
            // 上一步比较的结果 与 右节点比较，如果大于右节点，记录索引
            if p*2+2 < count && arr[min] > arr[p*2+2] {
                min = p*2+2
            }
            // 索引未改变，结束
            if min == p {
                break
            }
            // 交换，进行下一轮
            arr.swapAt(min, p)
            p = min
        }
    }
}

---

0x03 我的小作品

欢迎体验我的作品之一：小五笔86版应用
五笔学习好帮手
App Store 搜索即可~

---