最大堆与快排的例子

发布于 2024-01-12  375 次阅读


数组中第k个最大元素

给定整数数组 nums 和整数 k,请返回数组中第 k 个最大的元素。

请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

示例 1:

输入: [3,2,1,5,6,4], k = 2
输出: 5

示例 2:

输入: [3,2,3,1,2,4,5,5,6], k = 4
输出: 4

首先想到的当然是快排来解决,每次缩小查找范围,最后时间复杂度就是O(n)

官方代码:

class Solution {
    int quickselect(int[] nums, int l, int r, int k) {
        if (l == r) return nums[k];
        int x = nums[l], i = l - 1, j = r + 1;
        while (i < j) {
            do i++; while (nums[i] < x);
            do j--; while (nums[j] > x);
            if (i < j){
                int tmp = nums[i];
                nums[i] = nums[j];
                nums[j] = tmp;
            }
        }
        if (k <= j) return quickselect(nums, l, j, k);
        else return quickselect(nums, j + 1, r, k);
    }
    public int findKthLargest(int[] _nums, int k) {
        int n = _nums.length;
        return quickselect(_nums, 0, n - 1, n - k);
    }
}

代码很简洁,但细看逻辑还是有点复杂,这样写的快排有点难以理解,不过经过测试速度确实快,下面是我自己很好理解的快排:

class Solution {
    int quickselect(int[] nums, int left, int right, int k) {
        // 记录一下左右边界,方便递归
        int low = left;
        int high = right;
        int base = nums[left];
        while(left < right){
            // 从右边开始,寻找第一个小于基准的元素
            while (left < right && nums[right] >= base) {
                right--;
            }
            nums[left] = nums[right];

            // 从左边开始,寻找第一个大于基准的元素
            while (left < right && nums[left] <= base) {
                left++;
            }
            nums[right] = nums[left];
        }

        // 将基准元素放到最终的位置上
        nums[left] = base;

        if(left == k){
            return base;
        }
        // 递归处理左半部分
        else if(left>k){
            return quickselect(nums, low, left - 1,k);
        }
        // 递归处理右半部分
        else{
            return quickselect(nums, left + 1, high,k);
        }
    }
    public int findKthLargest(int[] _nums, int k) {
        int n = _nums.length;
        return quickselect(_nums, 0, n - 1, n - k);
    }
}

这个快排就比较清晰,认真看应该就能想通。

另外就是使用堆排序了,堆排序虽然时间复杂度是高了点O(nlogn),但这道题运行起来还蛮快的。

class Solution {
    public int findKthLargest(int[] nums, int k) {
        int heapSize = nums.length;
        buildMaxHeap(nums, heapSize);
        for (int i = nums.length - 1; i >= nums.length - k + 1; --i) {
            swap(nums, 0, i);
            --heapSize;
            maxHeapify(nums, 0, heapSize);
        }
        return nums[0];
    }

    public void buildMaxHeap(int[] a, int heapSize) {
        for (int i = heapSize / 2; i >= 0; --i) {
            maxHeapify(a, i, heapSize);
        } 
    }

    public void maxHeapify(int[] a, int i, int heapSize) {
        int l = i * 2 + 1, r = i * 2 + 2, largest = i;
        if (l < heapSize && a[l] > a[largest]) {
            largest = l;
        } 
        if (r < heapSize && a[r] > a[largest]) {
            largest = r;
        }
        if (largest != i) {
            swap(a, i, largest);
            maxHeapify(a, largest, heapSize);
        }
    }

    public void swap(int[] a, int i, int j) {
        int temp = a[i];
        a[i] = a[j];
        a[j] = temp;
    }
}

运行时间短的原因可能是因为测试用例中的k值较小。

堆排序还是挺重要的,这里使用数组实现的方法特别经典,可以认真看看。