leetcode215：数组中的第K个最大元素

题目链接

leetcode

题目描述

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

C++ 代码

#include <iostream>
#include <vector>
#include <queue>
using namespace std;

void printArray(vector<int>& nums) {
    for (const int& num : nums) {
        cout << num << " ";
    }
    cout << endl;
}

/*
    小根堆：堆的大小为 K，堆顶保存第 K 大的元素

    时间复杂度：O(NlogK)
        其中 N 是数组的长度，K 是要找的第 K 个最大元素的位置。
        遍历数组，需要 O(N) 的时间。每次插入操作和弹出操作，时间复杂度是O(logK)。
        因此，总的时间复杂度是O(NlogK)。
    空间复杂度：O(K)
        需要额外的空间，存储最大的 K 个元素。
*/
class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int, vector<int>, greater<int>> pq; // 小根堆

        for (const int& num : nums) {
            pq.push(num);
            if (pq.size() > k) {
                pq.pop();
            }
        }
        return pq.top();
    }
};

/*
    快速排序的选择方法，基于 Hoare 分区方案

    每次划分数组后，判断中轴值 pivot 在划分后的位置。 
    若位置正好等于 k ，则直接返回 pivot;
    否则，判断 pivot 在左边还是右边，继续寻找第 k 个最大元素。

    时间复杂度：O(n)
        递归时，每层时间复杂度为 O(n)，
        但不是都进入左右两部分递归，仅进入一侧递归在平均情况下数组长度会减半。
        故时间复杂度为 n + n/2 + n/4 + ... + 1 = O(n)。
    空间复杂度：O(logn)
        递归使用栈空间的大小是 O(logn)。
*/
class Solution_1 {
public:
    int findKthLargest(vector<int>& nums, int k) {
        return quickSelect(nums, 0, nums.size() - 1, k - 1);
    }
    
    int quickSelect(vector<int>& nums, int l, int r, int k) {
        if (l == r) {
            return nums[l];
        }

        // 随机选择基准值
        int pivotIndex = l + rand() % (r - l + 1);
        int pivot = nums[pivotIndex];
        // Hoare 分区，逆序排列
        int i = l, j = r;
        while (i <= j) {
            // 从左找 <= pivot
            while (i <= j && nums[i] > pivot) {
                i++;
            }
            // 从右找 >= pivot
            while (i <= j && nums[j] < pivot) {
                j--;
            }
            if (i <= j) {
                swap(nums[i], nums[j]);
                i++;
                j--;
            }
        }
        // nums[l..j] >= pivot
        // nums[i..r] <= pivot

        // 判断基准的位置
        if (k <= j) {
            return quickSelect(nums, l, j, k); // k 在左半部分
        } else if (k >= i) {
            return quickSelect(nums, i, r, k); // k 在右半部分
        } else {
            return nums[k]; // k 在中间，直接返回
        }
    }
};

/*
    思路同上，分区方案：荷兰国旗
*/
class Solution_2 {
public:
    int findKthLargest(vector<int>& nums, int k) {
        return quickSelect(nums, 0, nums.size() - 1, k - 1);
    }

    int quickSelect(vector<int>& nums, int l, int r, int k) {
        if (l == r) {
            return nums[l];
        }
    
        // 随机选择基准值
        int pivotIndex = l + rand() % (r - l + 1);
        int pivot = nums[pivotIndex];
        // 三路分区，逆序排列
        int low = l, mid = l, high = r; // 包含基准值自身在内的比较
        while (mid <= high) {
            if (nums[mid] > pivot) {
                swap(nums[mid], nums[low]);
                low++;
                mid++;
            } else if (nums[mid] < pivot) {
                swap(nums[mid], nums[high]);
                high--;
            } else {
                mid++;
            }
        }
        // nums[l..low-1] > pivot
        // nums[low..high] == pivot
        // nums[high+1..r] < pivot

        // 判断基准的位置
        if (k < low) {
            return quickSelect(nums, l, low - 1, k);
        } else if (k > high) {
            return quickSelect(nums, high + 1, r, k);
        } else {
            return nums[k];
        }
    }
};

int main() {
    Solution solution;
    vector<int> nums = {3, 2, 1, 5, 6, 4};
    printArray(nums);
    int k = 2;
    cout << "The " << k << "th largest element is " << solution.findKthLargest(nums, k) << endl;  // 应该输出 6

    return 0;
}

Golang 代码

package main

import (
	"container/heap"
	"fmt"
	"math/rand"
)

/*
   最小堆的使用：
       使用一个大小为 k 的最小堆来存储数组中的前 k 个最大元素。
       遍历数组时，如果堆的大小小于 k，直接将当前元素压入堆。
       如果堆的大小达到 k，且当前元素大于堆顶元素，则弹出堆顶元素并压入当前元素。
   堆的性质：
       最小堆的堆顶元素是最小的，因此堆顶元素即为当前的第 K 大元素。

   时间复杂度：O(nlogk)
       遍历数组的时间复杂度为 O(n)。
       每次插入或删除堆的操作的时间复杂度为 O(logk)。
       因此，总时间复杂度为 O(nlogk)。
   空间复杂度：O(k)
       使用了一个大小为 k 的堆，空间复杂度为 O(k)。
*/
// findKthLargest 返回数组中的第K个最大元素
func findKthLargest(nums []int, k int) int {
	h := &minHeap{}
	heap.Init(h)
	for _, num := range nums {
		if h.Len() < k {
			heap.Push(h, num)
		} else if num > (*h)[0] {
			heap.Pop(h)
			heap.Push(h, num)
		}
	}
	return (*h)[0]
}

// minHeap 是一个最小堆
type minHeap []int
func (h minHeap) Len() int           { return len(h) }
func (h minHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h minHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *minHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}
func (h *minHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]  // 获取最后一个元素
	*h = old[:n-1] // 移除最后一个元素
	return x
}

/*
	快速选择算法：
		将 partition 的逻辑直接嵌入到 quickSelect 函数中，减少了函数调用的开销。
		随机选择基准值并将其移动到右端，然后进行分区操作。
	分区操作：
		遍历数组，将小于基准值的元素移到左边，大于基准值的元素移到右边。
		最后将基准值放到正确的位置。
	递归选择：
		根据基准值的位置，决定在左半部分还是右半部分继续查找。

	时间复杂度：O(n)
		平均情况下，快速选择算法的时间复杂度为 O(n)。
		最坏情况下，时间复杂度为 O(n^2)，但通过随机化选择基准值，这种情况很少发生。
	空间复杂度：O(1)
		快速选择算法是原地操作，空间复杂度为 O(1)。
*/
// findKthLargest 使用快速选择算法找到数组中的第K个最大元素
func findKthLargest_1(nums []int, k int) int {
	// rand.Seed(time.Now().UnixNano())
	return quickSelect(nums, 0, len(nums)-1, len(nums)-k)
}

func quickSelect(nums []int, left, right, k int) int {
    if left == right {
        return nums[left]
    }

    // 随机选择基准值
	pivotIndex := rand.Intn(right-left+1) + left
    pivot := nums[pivotIndex]
	// Hoare 分区，逆序排列
    i, j := left, right
	for i <= j {
        for nums[i] < pivot { 
            i++
        }
        for nums[j] > pivot { 
            j--
        }
        if i <= j {
            nums[i], nums[j] = nums[j], nums[i]
            i++
            j--
        }
    }
	// nums[left..j] >= pivot
	// nums[i..right] <= pivot

	// 判断基准的位置
    if k <= j {
        return quickSelect(nums, left, j, k)
    } else if k >= i {
        return quickSelect(nums, i, right, k)
    } else {
        return nums[k]
    }
}

func main() {
	// 测试用例
	testCases := []struct {
		nums  []int
		k     int
		expected int
	}{
		{[]int{3, 2, 1, 5, 6, 4}, 2, 5},
		{[]int{3, 2, 3, 1, 2, 4, 5, 5, 6}, 4, 4},
	}

	for i, tc := range testCases {
		result := findKthLargest(tc.nums, tc.k)
		fmt.Printf("Test Case %d, Input: nums = %v, k = %d\n", i+1, tc.nums, tc.k)

		if result == tc.expected {
			fmt.Printf("Test Case %d, Output: %d, PASS\n", i+1, result)
		} else {
			fmt.Printf("Test Case %d, Output: %d, FAIL (Expected: %d)\n", i+1, result, tc.expected)
		}
	}
}