快速选择算法高效解决topk问题

快速选择算法高效解决topk问题

**你是否遇到过这样的业务问题,比如在数据库中根据某张表的某个字段来获取第K大(或小)的记录

• 也许你会使用order by 加limit的组合来返回该数据,但这样做如果数据量很大的话可能会需要很长的耗时。
• 或许你会选择返回所有数据,然后再排序处理返回排序后的len-K索引处即为结果,但这也需要O(nlogn)的时间复杂度
• 或者你会选择优先级队列,即将数据都加入优先级队列(用堆实现)中,然后pop出第k个数即为结果，时间也是O(nlogn)，建堆的时间代价是 O(n)，删除的总代价是O(klogn)，因为 k<n，故渐进时间复杂为O(n+klogn)=O(nlogn)。.

今天我推荐你使用快速选择来解决topK问题,可以在o(n)的时间内获取结果****

快速选择算法,可以在o(n)的时间内找到数组中第K大的元素

改进于快速排序算法

• 在快速排序中,我们每次续选择的基准值pivot,在一次快速排序过程后,在pivot所在索引处
  pivot索引左边的元素<=nums[pivot],右边的元素>=nums[pivot]
• 即pivot索引所在元素刚好是数组中的第nums.length-pivot大的元素,那么如果该索引刚好是我们需要的那个topK的索引,就无需再关心两边数据的顺序了
• nums.length-pivot==k,直接返回pivot索引处的元素 nums.length-pivot>k,向左递归该过程
  nums.length-pivot<k,向右递归该过程

以下是代码

public class QuickSelect {
    static Random random = new Random();

    public static void main(String[] args) {
        QuickSelect quickSelect = new QuickSelect();
        //假定是数据库返回的数据,用1亿条数据演示
        int[] nums = new int[100000000];
        int k = 3243252;//第k大的数,不超过数组长度都可以
        //用随机数填充数组
        getRandomNums(nums);
        //备用一份
        int[] nums2 = Arrays.copyOf(nums, nums.length);
        long l = System.currentTimeMillis();
        Arrays.sort(nums);
        long l1 = System.currentTimeMillis();
        System.out.println("排序耗时:" + (l1 - l) + "被选中的数:" + nums[nums.length - k]);

        long t = System.currentTimeMillis();
        int kthLargest = quickSelect.findKthLargest(nums2, k);
        long t1 = System.currentTimeMillis();
        System.out.println("快速选择耗时:" + (t1 - t) + "被选中的数:" + kthLargest);
    }

    private static void getRandomNums(int[] nums) {
        for (int i = 0; i < nums.length; i++) {
            nums[i] = random.nextInt(10000000);
        }
    }

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

    private int quickSelect(int[] nums, int l, int r, int index) {
        int i = randomPartition(nums, l, r);
        if (i == index) {
            return nums[i];
        }
        if (i < index) {
            return quickSelect(nums, i + 1, r, index);
        } else {
            return quickSelect(nums, l, i - 1, index);
        }
    }

    private int randomPartition(int[] nums, int l, int r) {
        //随机选取基准值索引
        int i = random.nextInt(r - l + 1) + l;
        //将基准值索引放在当前数组r索引处,这样就可以在[l,r-1]中搜索了，[l  .... r-1   pivot]
        swap(nums, i, r);
        //一次快排过程,使l...pivot-1都<=nums[pivot],pivot+1...r都>=nums[pivot]
        return partition(nums, l, r);
    }

    private int partition(int[] a, int l, int r) {
        int temp = a[r];
        int index = l - 1;//指向<=a[pivot]处索引
        for (int i = l; i < r; i++) {
            if (a[i] <= temp) {
                //将小的数放在pivot的左侧
                swap(a, ++index, i);
            }
        }
        //将pivot放到它应该在的位置
        swap(a, ++index, r);
        return index;
    }

    private void swap(int[] nums, int l, int r) {
        int temp = nums[l];
        nums[l] = nums[r];
        nums[r] = temp;
    }
}

运行几次看看效果,
k = 3243252

k = 8765432;

k = 12345678;

可以看到,在1亿的数据量情况下,快速选择几乎比排序快了10多倍