Find the TopK problem

Posted by programmer79 on Sat, 26 Oct 2019 21:06:58 +0200

Search for the largest / smallest k from a large number. This kind of problem is called TopK problem.

Data structure - maximum / small heap is usually used to solve the problem.

  • Find the front k, and use the minimum heap, that is, the element at the top of the heap is the minimum value in the heap.
  • The first k is small, and the maximum heap is used, that is, the element at the top of the heap is the maximum value in the heap.

If the value of the first k is large, enter the list=[12, 39, 3, 72, 56, 81, 15, 9, 103] and k=3, and output [103, 81, 72].
If the value of the first k is small, enter the list=[12, 39, 3, 72, 56, 81, 15, 9, 103] and k=3, and output [3, 9, 12].

Pre K large value

Train of thought:

Iteration list element:

  • 1. First put the first k elements to create a minimum heap;
  • 2. If the current element x is greater than the heap top element, remove the heap top element and merge it into team x;
  • 3. Finally, get the value in the minimum heap, which is topK4Max.

JAVA reference code

public static int[] topK4Max(int[] nums, int k) {// Priority queue, from small to large (default is also natural sort)
    Queue<Integer> minHeap = new PriorityQueue<>(Comparator.naturalOrder());
    // Build minimum heap
    for (int i = 0, len = nums.length; i < len; i++) {
        if (minHeap.size() < k) {
            minHeap.offer(nums[i]); // Join the team
            continue;
        }
        if (nums[i] > minHeap.peek()) { // Take team head element comparison, do not remove
            minHeap.poll(); // Take the team leader element and remove it
            minHeap.offer(nums[i]);
        }
    }
    // Traversal queue
    List<Integer> list = new ArrayList<>();
    Iterator<Integer> iterator = minHeap.iterator();
    while (iterator.hasNext()){
        list.add(iterator.next());
    }
    // Sort from large to small
    Collections.sort(list, (o1, o2) -> o1 < o2 ? 1 : o1 == o2 ? 0 : -1);
    // Convert to array
    return list.stream().mapToInt(Integer::valueOf).toArray();}

Pre K small value

Train of thought:

Iteration list element:

  • 1. Put the first k elements to create a maximum heap;
  • 2. If the current element x is smaller than the heap top element, remove the heap top element and merge it into team x;
  • 3. Finally, get the value in the largest pile, which is topK4Min.

JAVA reference code

public static int[] topK4Min(int[] nums, int k) {// Priority queue, sort from large to small
    Queue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
    // Build maximum heap
    for (int i = 0, len = nums.length; i < len; i++) {
        if (maxHeap.size() < k) {
            maxHeap.offer(nums[i]); // Join the team
            continue;
        }
        if (nums[i] < maxHeap.peek()) { // Take team head element for comparison, do not remove
            maxHeap.poll(); // Take the team leader element and remove it
            maxHeap.offer(nums[i]);
        }
    }
    // Traversal queue
    List<Integer> list = new ArrayList<>();
    Iterator<Integer> iterator = maxHeap.iterator();
    while (iterator.hasNext()){
        list.add(iterator.next());
    }
    // Sort from small to large
    Collections.sort(list);
    // Convert to array
    return list.stream().mapToInt(Integer::valueOf).toArray();
    
}

Test code

public static void main(String[] args) {
    int[] nums = new int[]{12, 39, 3, 72, 56, 81, 15, 9, 103};int[] maxK = topK4Max(nums, 3);
    System.out.println(JSON.toJSONString(maxK));int[] minK = topK4Min(nums, 3);
    System.out.println(JSON.toJSONString(minK));
}

More articles are welcome to scan the code and pay attention to the public account

Topics: Java JSON