Top K Largest Numbers

Question II (LI.545)

Implement a data structure that supports the following operations: 1. add(number) add the new number 2. topk() returns a list of top k numbers

Example

ds = new MyDS(3);
ds.add(3)
ds.add(10)
ds.topk()
>> return [10, 3]
ds.add(1000)
ds.add(-99)
s.topk()
>> return [1000, 10, 3]
ds.add(4)
ds.topk()
>> return [1000, 10, 4]

Approach

We want k largest so maintain a min heap (polling out the min) of size k.

Code

public class MyDS {
    private Queue<Integer> minHeap;
    private int maxSize;

    public MyDS(int k) {
        maxSize = k;
        minHeap = new PriorityQueue<Integer>(k);
    }

    public void add(int number) {
        if (minHeap.size() == maxSize) {
            if (number > minHeap.peek()) {
                minHeap.poll();
                minHeap.offer(number);
            }
        } else {
            minHeap.offer(number);
        }
    }

    public List<Integer> topk() {
        List<Integer> result = new ArrayList<>(minHeap);
        Collections.sort(result, Collections.reverseOrder());
        return result;
    }
}

Complexity

Time add(number) takes O(k) topk() takes O(klgk) (can be O(k) if not in sorted order)

Question I (LI.544)

Given an unsorted integer array, return the top k largest numbers.

Example

I: [3,10,1000,-99,4,100], k = 3
O: [1000, 100, 10]

Simple Heap

public int[] topk(int[] nums, int k) {
    // sanity check
    if (nums == null || k < 0 || nums.length < k) {
        return null;
    }
    int[] result = new int[k];
    // create a min heap of size k
    Queue<Integer> minHeap = new PriorityQueue<>(k);
    // add operation
    for (int num : nums) {
        if (minHeap.size() == k) {
            if (num > minHeap.peek()) {
                minHeap.poll();
                minHeap.offer(num);                    
            }
        } else {
            minHeap.offer(num);
        }
    }
    // priority queue to list
    for (int i = k - 1; i >= 0; i--) {
        result[i] = minHeap.poll();
    }
    return result;
}

Time O(nlgk) Space O(k)

Quick Select

O(logn) to find kth largest then O(n) to find the rest or O(nlogk) to sort the rest.

Last updated