1. 题目

描述

有一个整数数组,请你根据快速排序的思路,找出数组中第K大的数。
给定一个整数数组a,同时给定它的大小n和要找的K(1<=K<=n),请返回第K大的数(包括重复的元素,不用去重),保证答案存在。

示例1

  1. 输入:
  2. [1,3,5,2,2],5,3
  3. 复制
  4. 返回值:
  5. 2
  6. 示例2
  7. 复制

示例2

输入:
[10,10,9,9,8,7,5,6,4,3,4,2],12,3
复制
返回值:
9
说明:
去重后的第3大是8,但本题要求包含重复的元素,不用去重,所以输出9

2. 解(堆排序)

class Solution {
public:
    int findKth(vector<int> a, int n, int K) {
        // write code here
        buildHeap(a);
        int length = n;

        for (int i = 0; i < K; i++) {
            std::swap(a[0], a[length - 1 - i]);
            n--;
            siftRootDown(a, 0, n);
        }

        return a[length - K];
    }

    // 构建堆
    void buildHeap(vector<int>& a) {
        int length = a.size();
        if (length <= 1) return;

        int lastNonLeaf = (length - 2) >> 1;

        for (int i = lastNonLeaf; i >= 0; i--) {
            siftRootDown(a, i, length);
        }
    }

    // a: 实际的数组
    // root: 一棵子树的根节点
    // length: 堆的长度
    void siftRootDown(vector<int>& a, int root, int length){
        if (root >= length - 1)
            return;

        // 最小非叶子节点 (N - 1 - 1) / 2
        int minNonLeaf = (length - 2) >> 1;

        while (root <= minNonLeaf) {
            int leftChild = (root << 1) + 1;
            int rightChild = (root << 1) + 2;
            int maxIndex = root;

            // 左孩子存在   且   较大
            if (leftChild < length && a[maxIndex] < a[leftChild]) {
                maxIndex = leftChild;
            }
            // 右孩子存在   且   较大
            if (rightChild < length && a[maxIndex] < a[rightChild]) {
                maxIndex = rightChild;
            }

            if (maxIndex != root) {
                // 交换子树中最大的值和根节点的值
                std::swap(a[maxIndex], a[root]);
                // 检查子树是否满足堆序
                root = maxIndex;
                siftRootDown(a, root, length);
            }
            else {
                break;
            }
        }
    }
};