求第K大算法总结

Posted 2023-02-22 馥钰

搜索了一下共发现两种解法，一种是借助一个额外的堆空间，一种是利用快排的方式根据位置来查找！

解法一：借助辅助堆

建立一个k个元素的小顶堆，首先将数组中前k个元素放入堆中，此时堆顶元素为第k大的元素，后面继续遍历数组，比较数组中元素与堆顶元素值，当数组中元素大于堆顶元素时，将堆顶元素弹出，新元素入堆，这样最终堆顶元素即为第k大。

可以直接利用Java中优先级队列，这里传入一个比较器来构造小顶堆

import org.junit.Test;

import java.util.Comparator;
import java.util.PriorityQueue;

public class TopKTest 

    @Test
    public void topKQueue() 
        int[] arr = 1, 2, 3, 4, 5, 6, 7, 8, 8, 9, 2, 2, 23, 45, 78, 23, 45, 56;
        int topK = topK(arr, 3);
        System.out.println(topK);
    


    public static int topK(int[] arr, int k) 
        //建立小根堆
        PriorityQueue<Integer> q = new PriorityQueue<>(k, new Comparator<Integer>() 
            @Override
            public int compare(Integer o1, Integer o2) 
                //元素升序排序
                return o1 - o2;
            
        );
        for (int i = 0; i < k; i++) 
            q.add(arr[i]);
        
        for (int i = k; i < arr.length; i++) 
            //比堆顶元素大，则移除堆顶元素，加入堆
            if (arr[i] > q.peek()) 
                q.poll();
                q.add(arr[i]);
            
        
        return q.peek();
    

解法二：利用快排

利用快速排序中partition函数，当partition返回元素坐标q==k-1时，说明q前面已经有k-1个元素比q位置元素小，则arr[q]即为第k大元素。如果q>k-1 说明第k大元素在前半部分，否则在后半部分。

import org.junit.Test;

public class FindTopKTest 

    @Test
    public void topKTest() 
        int[] arr = 1, 2, 3, 4, 5, 6, 7, 8, 8, 9, 2, 2, 23, 45, 78, 23, 45, 56;
        int topK3 = topK(arr, 3);
        System.out.println(topK3);
    


    public static int topK(int[] nums, int k) 
        return topK(nums, k, 0, nums.length - 1);
    
    
    public static int topK(int[] nums, int k, int start, int end) 
        //第k大个元素
        int low = start;
        int high = end;
        int mid = nums[low];
        while (low < high) 
            //找到后半部分[x - high]（low < x < high）小于mid，移到low位置
            while (low < high && nums[high] <= mid) high--;
            if (low < high) nums[low] = nums[high];

            //找到前半部分[low - x]（low < x < high）大于等于mid，移到high位置
            while (low < high && nums[low] >= mid) low++;
            if (low < high) nums[high] = nums[low];
        
        nums[low] = mid;
        if (low == k - 1) 
            return mid;
         else if (low > k - 1) 
            //第k大位于前半部分
            return topK(nums, k, start, low - 1);
         else 
            //第K大位于后半部分
            return topK(nums, k, low + 1, end);
        
    

参考

面试题：从n个数中找出第K大的数_orangefly0214的博客-CSDN博客_第k大的数