海量数据 Top K 问题的解决方案

Top K是很常见的一种问题，是指在N个数的无序序列中找出最大的K个数，而其中的N往往都特别大，对于这种问题，最容易想到的办法当然就是先对其进行排序，然后直接取出最大的K个元素就行了，但是这种方法往往是不可靠的，不仅时间效率低而且空间开销大，排序是对所有数都要进行排序，而实际上，这类问题只关心最大的K个数，并不关心序列是否有序，因此，排序实际上是浪费了的很多资源都是没必要的。

题目：

输入 n 个整数，找出其中最大的 k 个数。例如输入4、5、1、6、2、7、3、8 这8个数字，则最大的4个数字是5、6、7、8。

解法一：基于快排的O(n)的算法
如果基于数组的第 k 个数字来调整，使得比第 k 个数字小的所有数字都位于数组的左边，比第 k 个数字大的所有数字都位于数组的右边。这样调整之后，位于数组中左边的 k 个数字就是最小的 k 个数字（这 k 个数字不一定是排序的）。

public class LeastK {
    public static void getLeastNumbers(int[] input, int[] output) {
        if (input == null || output == null || output.length <= 0 || input.length < output.length) {
            throw new IllegalArgumentException("Invalid args");
        }
        int start = 0;
        int end = input.length - 1;
        int index = partition(input, start, end); // 切分后左子数组的长度
        int target = output.length - 1; // K-1
 
        // 若切分后左子数组长度不等于K
        while (index != target) { // 若切分后左子数组长度小于K，那么继续切分右子数组，否则继续切分左子数组
            if (index < target) {
                start = index + 1;
            } else {
                end = index - 1;
            }
            index = partition(input, start, end);
        }
 
        System.arraycopy(input, 0, output, 0, output.length);
    }
 
    private static int partition(int arr[], int left, int right) {
        int i = left;
        int j = right + 1;
        int pivot = arr[left];
        while (true) { // 找到左边大于pivot的数据，或者走到了最右边仍然没有找到比pivot大的数据
            while (i < right && arr[++i] < pivot) { // 求最大的k个数时，arr[++i] > pivot
                if (i == right) {
                    break;
                }
            } // 找到右边小于pivot的数据，或者走到了最左边仍然没有找到比pivot小的数据
            while (j > left && arr[--j] > pivot) { // 求最大的k个数时，arr[--j] < pivot
                if (j == left) {
                    break;
                }
            } // 左指针和右指针重叠或相遇，结束循环
            if (i >= j) {
                break;
            } // 交换左边大的和右边小的数据
            swap(arr, i, j);
        } // 此时的 a[j] <= pivot，交换之
        swap(arr, left, j);
        return j;
    }
 
    private static void swap(int[] arr, int i, int j) {
        int tmp = arr[i];
        arr[i] = arr[j];
        arr[j] = tmp;
    }
 
}

采用上面的思路是有限制的，比如需要修改输入的数组，因为函数 Partition 会调整数组中的顺序，当然了，这个问题完全可以通过事先拷贝一份新数组来解决。值得说明的是，这种思路是不适合处理海量数据的。若是遇到海量数据求最小的 k 个数的问题，可以使用下面的解法。

解法二：最小堆法

利用最小堆的思想，先读取前k个元素，建立一个最小堆。然后将剩余的所有元素依次与堆顶元素进行比较，如果大于堆顶元素，则堆顶弹出，否则，压入下一个数组元素继续比较，只要维护大小为k的堆就可以了，此方法适合处理海量数据，时间复杂度为O(nlogk)。java的PriorityQueue可以实现最小堆。

import java.util.Arrays;
import java.util.PriorityQueue;
import java.util.Random;
 
public class TopK {
 
    public static void main(String[] args) {
        int number = 100000000;// 一亿个数
        int maxnum = 1000000000;// 随机数最大值
        int topnum = 100;// 取最大的多少个
 
        int[] nums = new int[100000000];
        Random random = new Random();
        for (int i = 0; i < number; i++) {
            int s = Math.abs(random.nextInt(maxnum));
            nums[i] = s;
        }
 
        Integer[] res = TopK.getLargestNumbers(nums, topnum);
        System.out.println(Arrays.toString(res));
 
    }
 
    public static Integer[] getLargestNumbers(int[] nums, int k) {
        PriorityQueue<Integer> minQueue = new PriorityQueue<>(k); // 默认自然排序
        for (int num : nums) {
            if (minQueue.size() < k || num > minQueue.peek()) { // peek()：返回队列头部的值，也就是队列最小值
                // 插入元素
                minQueue.offer(num);
            }
            if (minQueue.size() > k) { // 删除队列头部
                minQueue.poll();
            }
        }
        return minQueue.toArray(new Integer[0]);
    }
 
}

也可以将上述代码改为手动实现最小堆的方式，代码如下：

 
import java.util.Arrays;
import java.util.Date;
import java.util.Random;
 
public class Top100 {
    
    public static void main(String[] args) {
        find();
    }
    public static void find( ) {//
        int number = 100000000;// 一亿个数
        int maxnum = 1000000000;// 随机数最大值
        int i = 0;
        int topnum = 100;// 取最大的多少个
     
        Date startTime = new Date();
        
        Random random = new Random();
        int[] top = new int[topnum];
        for (i = 0; i < topnum; i++) {
            top[i] = Math.abs(random.nextInt(maxnum));//设置为随机数
//            top[i] = getNum(i);
        }
 
        buildHeap(top, 0, top.length);// 构建最小堆， top[0]为最小元素
        for (i = topnum; i < number; i++) {
 
            int currentNumber2 = Math.abs(random.nextInt(maxnum));//设置为随机数
//            int currentNumber2 = getNum(i);
            // 大于 top[0]则交换currentNumber2  重构最小堆
            if (top[0] < currentNumber2) {
                top[0] = currentNumber2;
                shift(top, 0, top.length, 0); // 构建最小堆 top[0]为最小元素
            }
        }
        System.out.println(Arrays.toString(top));
        sort(top);
        System.out.println(Arrays.toString(top));
        
        Date endTime = new Date();
        System.out.println("用了"+(endTime.getTime() - startTime.getTime())+"毫秒");
 
    }
    
    public static int getNum(int i){
        return i;
    }
 
 
    //构造排序数组
    public static void buildHeap(int[] array, int from, int len) {
        int pos = (len - 1) / 2;
        for (int i = pos; i >= 0; i--) {
            shift(array, from, len, i);
        }
    }
 
    /**
     * @param array top数组
     * @param from 开始
     * @param len 数组长度
     * @param pos 当前节点index
     * */
    public static void shift(int[] array, int from, int len, int pos) {
        // 保存该节点的值 
        int tmp = array[from + pos];
 
        int index = pos * 2 + 1;// 得到当前pos节点的左节点
        while (index < len)//  存在左节点
        {
            if (index + 1 < len
                    && array[from + index] > array[from + index + 1])// 如果存在右节点
            {
                // 如果右边节点比左边节点小，就和右边的比较
                index += 1;
            }
             
            if (tmp > array[from + index]) {
                array[from + pos] = array[from + index];
                pos = index;
                index = pos * 2 + 1;
            } else {
                break;
            }
        }
        // 最终全部置换完毕后 ，把临时变量赋给最后的节点
        array[from + pos] = tmp;
    }
    
    public static void sort(int[] array){  
        for(int i = 0; i < array.length - 1; i++){  
            //当前值当作最小值  
            int min = array[i];  
            for(int j = i+1; j < array.length; j++){  
                if(min>array[j]){  
                    //如果后面有比min值还小的就交换  
                    min = array[j];  
                    array[j] = array[i];  
                    array[i] = min;  
                }  
            }  
        }  
    }
 
}

若是遇到此类求海量数据中最小的 k 个数的问题，只需改用最大堆即可。构建最大堆，需要重写compare方法。

import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Random;
 
public class TopKK {
 
    public static void main(String[] args) {
        int number = 100000000;// 一亿个数
        int maxnum = 1000000000;// 随机数最大值
        int topnum = 100;// 取最大的多少个
 
        int[] nums = new int[100000000];
        Random random = new Random();
        for (int i = 0; i < number; i++) {
            int s = Math.abs(random.nextInt(maxnum));
            nums[i] = s;
        }
 
        Integer[] res = TopKK.getLeastNumbers(nums, topnum);
        System.out.println(Arrays.toString(res));
 
    }
 
    public static Integer[] getLeastNumbers(int[] nums, int k) { // 默认自然排序，需手动转为降序
        PriorityQueue<Integer> maxQueue = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                //if (o1 > o2) {
                //    return -1;
                //} else if (o1 < o2) {
                //    return 1;
                //}
                //return 0;
                return o2.compareTo(o1);
            }
        });
        for (int num : nums) {
            if (maxQueue.size() < k || num < maxQueue.peek()) { // peek()：返回队列头部的值，也就是队列最大值
                // 插入元素
                maxQueue.offer(num);
            }
            if (maxQueue.size() > k) { // 删除队列头部
                maxQueue.poll();
            }
        }
        return maxQueue.toArray(new Integer[0]);
    }
 
}

本文已通过「原本」原创作品认证，转载请注明文章出处及链接。

海量数据 Top K 问题的解决方案

发表评论