Smallest K Numbers

1. Problem Description

Given n integers, find the smallest K numbers. For example, among 4,5,1,6,2,7,3,8, the smallest four numbers are 1,2,3,4.

2. Approach

• Quick sort
• Max heap: to find the smallest k numbers, use a max heap whose root is the largest among the k retained numbers

3. Code

3.1. Sorting

// Time: O(NlogN)
func GetLeastNumbers_Solution(input []int, k int) []int {
	if len(input) == 0 {
		return nil
	}

	if k <= 0 || k > len(input) {
		return nil
	}

	quickSort(input, 0, len(input)-1)
	kNumbers := make([]int, 0)
	for i := 0; i < k; i++ {
		kNumbers = append(kNumbers, input[i])
	}

	return kNumbers
}

func quickSort(data []int, left int, right int) {
	if left >= right {
		return
	}

	p := partition(data, left, right)
	quickSort(data, left, p-1)
	quickSort(data, p+1, right)
}

func partition(data []int, left int, right int) int {
	// Generate a random index in [l, r]
	p := left + randInt(0, right-left)
	swap(data, left, p)

	//data[left+1...j] <= v; data[j+1...i] >= v
	i := left + 1
	j := right
	for {
		for i <= j && data[i] < data[left] {
			i++
		}
		for j >= i && data[j] > data[left] {
			j--
		}
		if i >= j {
			break
		}

		swap(data, i, j)
		i++
		j--
	}

	swap(data, left, j)
	return j
}

func swap(data []int, i int, j int) {
	tmp := data[i]
	data[i] = data[j]
	data[j] = tmp
}

func randInt(min, max int) int {
	if min >= max || min == 0 || max == 0 {
		return max
	}
	return rand.Intn(max-min) + min
}

3.2. Heap

• java

public class 最小的K个数
{
    public ArrayList<Integer> GetLeastNumbers_Solution(int[] input, int k)
    {
        ArrayList<Integer> res = new ArrayList<>();
        // Check parameters
        if (input == null || input.length == 0 || k <= 0 || k > input.length)
        {
            return res;
        }
        // The smallest k numbers can be maintained with a max heap whose root is the largest among those k numbers
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(k, new Comparator<Integer>()
        {
            @Override
            public int compare(Integer o1, Integer o2)
            {
                return o2 - o1;
            }
        });
        // Traverse the max heap and insert values into the list
        for (int val : input)
        {
            // If the count has not reached k, insert into the max heap
            if (maxHeap.size() < k)
            {
                maxHeap.offer(val);
            }
            else
            {
                // If val is smaller than the heap root, remove the root and insert this value
                if (val < maxHeap.peek())
                {
                    maxHeap.poll();
                    maxHeap.offer(val);
                }
            }
        }

        for (Integer val : maxHeap)
        {
            res.add(val);
        }

        return res;
    }
}

• go

import (
	"container/heap"
	"sort"
)

type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func (h *IntHeap) Push(v interface{}) {
	*h = append(*h, v.(int))
}

func GetLeastNumbers_Solution2(input []int, k int) []int {
	// write code here
	if k <= 0 || k > len(input) {
		return []int{}
	}
	//     sort.Ints(input)
	//     return input[:k]
	h := IntHeap{}
	heap.Init(&h)
	for _, v := range input {
		if h.Len() < k {
			heap.Push(&h, v)
		} else if h.Len() == k {
			temp := heap.Pop(&h).(int)
			if temp > v {
				heap.Push(&h, v)
			} else {
				heap.Push(&h, temp)
			}
		}
	}
	sort.Ints(h)
	return h
}

4. References

• Smallest K Numbers - NowCoder
• Interview Question 17.14. Smallest K Numbers - LeetCode