Kth Smallest Number in M Sorted Lists
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)
.