Kth Smallest Number in M Sorted Lists

Ethan Davis
Data Structures and Algorithms DSA
3 min readApr 2, 2024
Data Structures and Algorithms

Statement

Given m number of sorted lists in ascending order and an integer k, find the kth smallest number among all given lists.

There can be repeated values in the lists, but each value is considered unique, and counts towards finding the kth smallest number.

If k is greater than the total number of values, then return the greatest value, and if there’s not values, then return 0.

Constraints

  • 1 ≤ m ≤ 50
  • 0 ≤ lists[i].length() ≤ 50
  • -10⁹ ≤ lists[i][j] ≤ 10⁹
  • 1 ≤ k ≤ 10⁹

Solution

"""
production algorithm
"""

from data_structures.heap.heap import Heap


class Solution:
def k_smallest_number(self, lists, k):
"""
time complexity O(klogn)
space complexity O(n)
"""
# add first number of each list to heap
smallest = Heap()
pointers = [0 for _ in range(len(lists))]
for i in range(len(lists)):
if 0 < len(lists[i]):
smallest.push(MinHeapData(lists[i][pointers[i]], i))
result = None

for i in range(k):
if smallest.empty():
break

# add number of list from heap to result
number, index = smallest.pop()
result = number

# add next number of list to heap
pointers[index] = pointers[index] + 1
if pointers[index] < len(lists[index]):
smallest.push(MinHeapData(
lists[index][pointers[index]], index))

if result is None:
return 0
return result


class MinHeapData:
def __init__(self, value, index):
self.value = value
self.index = index

def __lt__(self, other):
return self.value < other.value

def __iter__(self):
yield self.value
yield self.index

def __len__(self):
return 2
"""
unit tests
"""

from unittest import TestCase
from algorithms.k_way_merge.kth_smallest_number.solution import Solution


class SolutionTestCase(TestCase):
def test_all_empty_lists(self):
# Given
lists = [[], [], [], [], []]
k = 7
solution = Solution()

# When
actual = solution.k_smallest_number(lists, k)

# Then
expected = 0
self.assertEqual(actual, expected)

def test_some_empty_lists(self):
# Given
lists = [[], [1, 3, 5, 6, 9], [0, 2, 4, 6, 8], [], [-3, -2, -1, 0, 1]]
k = 7
solution = Solution()

# When
actual = solution.k_smallest_number(lists, k)

# Then
expected = 1
self.assertEqual(actual, expected)

def test_no_empty_lists(self):
# Given
lists = [[53, 55, 57, 59], [35, 45, 55, 65], [
49, 59, 69, 79], [20, 40, 60, 80], [101, 201, 301, 401]]
k = 7
solution = Solution()

# When
actual = solution.k_smallest_number(lists, k)

# Then
expected = 55
self.assertEqual(actual, expected)

def test_total_less_than_k(self):
# Given
lists = [[53, 55, 57, 59], [35, 45, 55, 65], [
49, 59, 69, 79], [20, 40, 60, 80], [101, 201, 301, 401]]
k = 30
solution = Solution()

# When
actual = solution.k_smallest_number(lists, k)

# Then
expected = 401
self.assertEqual(actual, expected)

Strategy

K-way Merge.

Explanation

First, push the first number of each of the n sorted lists to a min heap. Then, for at most k iterations, pop the smallest number from the heap. When a number of a list is popped from the heap, then push the next number of that list to the heap. After k iterations, the kth smallest number will be found.

Time Complexity

Let m be the number of sorted lists. At most, one number of each list will be in the min heap at any given time. Then, at most, there are k iterations of popping from and pushing to the min heap, in order to find the kth smallest number of m sorted lists. Therefore, the time complexity of the algorithm is O(klogm).

Space Complexity

The auxiliary space complexity of the algorithm depends on the size of the min heap used and is at most size m. Therefore, the auxiliary space complexity of the algorithm is O(m).

Links

--

--