Solution: Kth Smallest Number in M Sorted Lists

Let's solve the Kth Smallest Number in M Sorted Lists problem using the K-Way Merge pattern.

We'll cover the following

Statement

Given an mm number of sorted lists in ascending order and an integer, k, find the kthk^{th} smallest number among all the given lists.

Although there can be repeating values in the lists, each element is considered unique and, therefore, contributes to calculating the kthk^{th} smallest element.

If k is greater than the total number of elements in the input lists, return the greatest element from all the lists and if there are no elements in the input lists, return 0.

Constraints:

  • 1≤1\leq m ≤50\leq50
  • 0≤0\leq list[i].length ≤50\leq 50
  • −109≤-10^9\leq list[i][j] ≤109\leq 10^9
  • 1≤1\leq k ≤109\leq 10^9

Solution

So far, you’ve probably brainstormed some approaches and have an idea of how to solve this problem. Let’s explore some of these approaches and figure out which one to follow based on considerations such as time complexity and any implementation constraints.

Naive approach

A naive approach to this problem can be to merge all the lists into a single list and then sort it to find the kthk^{th} smallest number from it.

The time complexity of the naive approach is O(nlog⁡n)O(n \log n), where nn is the total number of elements present in all the lists. Since sorting takes O(nlog⁡n)O(n \log n), while searching up to kthk^{th} number on the sorted list costs us O(k)O(k).

Optimized approach using k-way merge

This approach uses a min-heap to combine the lists into a single, sorted list in ascending order. Initially, the first element of each list is added to the min-heap to ensure that we always have the smallest available element at the top of the heap. Then, an element is removed from the heap, and if the removed element has a next element in its original list, that next element is added to the heap. This process is repeated until kk elements have been removed. The kthk^{th} element removed is the kthk^{th} smallest element across all the lists.

Note: We’ll implement the heap using the PriorityQueue class.

The slides below illustrate how we would like the algorithm to run.

Let’s look at the code for this solution below:

import java.util.*;
 
class KSmallestNumber {
    public static int kSmallestNumber(List<List<Integer>> lists, int k) {
        int listLength = lists.size();
        PriorityQueue<int[]> kthSmallest = new PriorityQueue<>((a, b) -> a[0] - b[0]);
        for (int index = 0; index < listLength; index++) {
            if (lists.get(index).size() == 0) {
                continue;
            } else {
                kthSmallest.offer(new int[] {lists.get(index).get(0), index, 0});
            }
        }
        int numbersChecked = 0, smallestNumber = 0;
        while (!kthSmallest.isEmpty()) {  
            int[] smallest = kthSmallest.poll();
            smallestNumber = smallest[0];
            int listIndex = smallest[1];
            int numIndex = smallest[2];
            numbersChecked++;
            if (numbersChecked == k) {
                break;
            }
            if (numIndex + 1 < lists.get(smallest[1]).size()) {
                kthSmallest.offer(new int[] {lists.get(listIndex).get(numIndex + 1), listIndex, numIndex + 1});
            }
        }
        return smallestNumber;
    }
    // Driver code
    public static void main(String[] args) {
        
        List<List<List<Integer>>> lists = Arrays.asList(
            Arrays.asList(
                Arrays.asList(2, 6, 8),
                Arrays.asList(3, 6, 10),
                Arrays.asList(5, 8, 11)
            ),
            Arrays.asList(
                Arrays.asList(1, 2, 3),
                Arrays.asList(4, 5),
                Arrays.asList(6, 7, 8, 15),
                Arrays.asList(10, 11, 12, 13),
                Arrays.asList(5, 10)
            ),
            Arrays.asList(
                Arrays.asList(),
                Arrays.asList(),
                Arrays.asList()
            ),
            Arrays.asList(
                Arrays.asList(1, 1, 3, 8),
                Arrays.asList(5, 5, 7, 9),
                Arrays.asList(3, 5, 8, 12)
            ),
            Arrays.asList(
                Arrays.asList(5, 8, 9, 17),
                Arrays.asList(),
                Arrays.asList(8, 17, 23, 24)
            )
        );
        int[] k = {5, 50, 7, 4, 8};
        for (int i = 0; i < k.length; i++) {
            System.out.println(i + 1 + ".\t Input lists: " + lists.get(i) +
                "\n\t K = " + k[i] +
                "\n\t " + k[i] + "th smallest number from the given lists is: " +
                kSmallestNumber(lists.get(i), k[i]));
            System.out.println(new String(new char[100]).replace('\0', '-'));
        }
    }
}
Kth Smallest Number in M Sorted Lists

Solution summary

We can summarize our solution in the following steps.

  1. Push the first element of each list in the min-heap.

  2. Pop the top element of the min-heap, and keep track of the list index and the element index in the list. Now, push the next element of the popped element from the list if the element index is not the last index in the list.

  3. Repeat step 2 until we have popped kk elements from the heap. If we have popped kk elements, then the kthk^{th} element popped in this process is the kthk^{th} smallest element.

Time complexity

The time complexity of the first step is

  • O(m)O(m) for iterating the first element of each list, where mm is the number of lists.

  • The cost of pushing mm elements in the heap is as follows:

log⁡1+log⁡2+log⁡3+⋯+log⁡m=log⁡(1∗2∗3∗⋯∗m)=log⁡(m!)\log 1 + \log 2 + \log 3 + \dots + \log m = \log (1*2*3* \dots *m) = \log(m!)

As per Stirling’s approximation, O(log⁡m!)≈O(mlog⁡m)O(\log m!) \approx O(m \log m).

So, the time complexity of the first loop is O(mlog⁡m)O(m \log m).

  • In the while loop, we pop and push on to the heap kk number of times until we find the kthk^{th} smallest number. So, the time complexity of this step is O(klog⁡m)O(k \log m).

So, the total time complexity of this solution is

O(mlog⁡m+klog⁡m)=O((m+k)log⁡m)O(m \log m + k \log m) = O((m+k) \log m)

Space complexity

The space complexity is O(m)O(m), where mm is the total number of elements in the heap.

Level up your interview prep. Join Educative to access 70+ hands-on prep courses.