百万级数据量下找到数组中的第K个最大元素

相信看过我这篇文章的朋友可能会对这些不同时间复杂度的实现有着好奇，到底它们的差别会有多大呢？今天我们自己造数据，运行一下看看是什么效果。

问题：

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

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

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

造数据

造数据我用Python比较方便，下面的代码是Python3.8环境运行的：

Python 复制代码

import sys
import random

def main():
    if len(sys.argv) < 2:
        print("请提供要生成的数字数量作为命令行参数。")
        sys.exit(1)
    try:
        n = int(sys.argv[1])
    except ValueError:
        print("参数必须为整数。")
        sys.exit(1)

    # 生成包含正负数的大范围随机数
    nums = [str(random.randint(-10**5, 10**5)) for _ in range(n)]
    print(' '.join(nums))

if __name__ == "__main__":
    main()

给上面的代码命名为generate.py，然后在命令行使用下面的方式就能造出百万数据了：

bash 复制代码

python generate.py 1000000 > data.txt

读数据，算法比较

只要打开Ide，创建一个类就好了：

Java 复制代码

package hot100;


import java.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.util.Arrays;
import java.util.PriorityQueue;

public class KthLargest {
    public static void main(String[] args) {
        String filePath = "C:\Users\xqkon\Documents\Work\Java Projects\leetcode\src\hot100\data.txt";
        int k = 1000; // 可根据需要修改或通过命令行参数传入

        try (BufferedReader br = new BufferedReader(new FileReader(filePath))) {
            String line = br.readLine();
            String[] parts = line.split("\s+");
            int[] nums = new int[parts.length];

            // 将字符串转换为整数数组
            for (int i = 0; i < parts.length; i++) {
                nums[i] = Integer.parseInt(parts[i]);
            }

            // 性能测试
            long startTime = System.nanoTime();
            int result = findKthLargestForce(nums, k);
            long endTime = System.nanoTime();

            // 输出结果和时间（毫秒）
            System.out.printf("第%d大元素: %d%n", k, result);
            System.out.printf("耗时: %.3f 毫秒%n", (endTime - startTime) / 1e6);

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private static int findKthLargest(int[] nums, int k) {
        return quickSelect(nums, 0, nums.length - 1, nums.length - k);
    }

    // 快速选择算法
    private static int quickSelect(int[] nums, int left, int right, int k) {
        if (left == right) return nums[left];

        int pivotIndex = partition(nums, left, right);

        if (k == pivotIndex) return nums[k];
        else if (k < pivotIndex) return quickSelect(nums, left, pivotIndex - 1, k);
        else return quickSelect(nums, pivotIndex + 1, right, k);
    }

    private static int partition(int[] nums, int left, int right) {
        int pivot = nums[right];
        int i = left;
        for (int j = left; j < right; j++) {
            if (nums[j] <= pivot) {
                swap(nums, i, j);
                i++;
            }
        }
        swap(nums, i, right);
        return i;
    }

    private static void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }

    // 使用类库里的最小堆
    public static int findKthLargestPQ(int[] nums, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<Integer>();  // also can add param: (a, b) -> a - b, still 最小堆 (min - heap)
        for (int i : nums) {
            if (heap.size() < k) {
                heap.offer(i);
            } else if (i > heap.peek()) {
                heap.poll();
                heap.offer(i);
            }
        }
        return heap.peek();
    }

    // 暴力解法
    public static int findKthLargestForce(int[] nums, int k) {
        Arrays.sort(nums);
        return nums[nums.length-k];
    }
}

接下来是三种算法的比较(k = 100)：

暴力方法 O(nlogn)：

bash 复制代码

第1000大元素: 99800
耗时: 91.194 毫秒

最小堆 O(nlogk)：

bash 复制代码

第1000大元素: 99800
耗时: 41.477 毫秒

快速选择排序平均O(n)：

bash 复制代码

第1000大元素: 99800
耗时: 7.565 毫秒

由此可见O(n)级别算法的厉害之处，当数据量来到百万级别时间直接差了一个量级。想想看，如果是互联网公司，动辄上亿的数据量呢？可想而知，将心比心，大家为什么要考算法。