Skip to content

17.14. Smallest K

Description

Design an algorithm to find the smallest K numbers in an array.

Example: 

Input:  arr = [1,3,5,7,2,4,6,8], k = 4

Output:  [1,2,3,4]

Note:

  • 0 <= len(arr) <= 100000
  • 0 <= k <= min(100000, len(arr))

Solutions

Solution 1

Thinking

The \(k\) smallest values, in any order. A full sort plus a prefix is optimal when \(k\) is close to \(n\).

sorted(arr)[:k] is short and has a good constant in Python.

The statement allows any order, so no extra heap is required here.

1
2
3
class Solution:
    def smallestK(self, arr: List[int], k: int) -> List[int]:
        return sorted(arr)[:k]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
class Solution {
    public int[] smallestK(int[] arr, int k) {
        Arrays.sort(arr);
        int[] ans = new int[k];
        for (int i = 0; i < k; ++i) {
            ans[i] = arr[i];
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        sort(arr.begin(), arr.end());
        vector<int> ans(k);
        for (int i = 0; i < k; ++i) {
            ans[i] = arr[i];
        }
        return ans;
    }
};

1
2
3
4
5
6
7
8
func smallestK(arr []int, k int) []int {
    sort.Ints(arr)
    ans := make([]int, k)
    for i, v := range arr[:k] {
        ans[i] = v
    }
    return ans
}

1
2
3
4
5
6
7
class Solution {
    func smallestK(_ arr: [Int], _ k: Int) -> [Int] {
        guard k > 0 else { return [] }
        let sortedArray = arr.sorted()
        return Array(sortedArray.prefix(k))
    }
}

Solution 2

Thinking

A full sort wastes work when \(k\ll n\).

A max-heap of size \(k\) keeps the current \(k\) smallest while scanning the rest, in \(O(n\log k)\).

1
2
3
4
5
6
7
8
class Solution:
    def smallestK(self, arr: List[int], k: int) -> List[int]:
        h = []
        for v in arr:
            heappush(h, -v)
            if len(h) > k:
                heappop(h)
        return [-v for v in h]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution {
    public int[] smallestK(int[] arr, int k) {
        PriorityQueue<Integer> q = new PriorityQueue<>((a, b) -> b - a);
        for (int v : arr) {
            q.offer(v);
            if (q.size() > k) {
                q.poll();
            }
        }
        int[] ans = new int[k];
        int i = 0;
        while (!q.isEmpty()) {
            ans[i++] = q.poll();
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        priority_queue<int> q;
        for (int& v : arr) {
            q.push(v);
            if (q.size() > k) {
                q.pop();
            }
        }
        vector<int> ans;
        while (q.size()) {
            ans.push_back(q.top());
            q.pop();
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
func smallestK(arr []int, k int) []int {
    q := hp{}
    for _, v := range arr {
        heap.Push(&q, v)
        if q.Len() > k {
            heap.Pop(&q)
        }
    }
    ans := make([]int, k)
    for i := range ans {
        ans[i] = heap.Pop(&q).(int)
    }
    return ans
}

type hp struct{ sort.IntSlice }

func (h hp) Less(i, j int) bool { return h.IntSlice[i] > h.IntSlice[j] }
func (h *hp) Push(v any)        { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() any {
    a := h.IntSlice
    v := a[len(a)-1]
    h.IntSlice = a[:len(a)-1]
    return v
}

Comments