Katkat
Back

Return the K Smallest Values from Sorted Arrays

eBaySoftware Engineer · Technical Screen
MediumSoftware EngineerTechnical ScreenCoding & Algorithms

Last updated

Tests ability to efficiently merge multiple sorted streams to extract the k smallest elements using a min-heap or frontier-based approach. Evaluates handling of empty arrays, duplicates, and computational complexity dependent on the number of arrays rather than total elements.

Problem Statement

You are given a list of m integer arrays arrays, where each individual array is sorted in non-decreasing order. You are also given an integer k.

Write a function to return the k smallest elements collected across all arrays, sorted in non-decreasing order. If an element appears multiple times across or within the arrays, each occurrence should be preserved. If the total number of elements across all arrays is less than k, return all elements sorted in non-decreasing order.

Your solution should aim for an efficiency that depends on m (the number of arrays) and k, rather than concatenating and sorting all input elements.

Function Signature

def k_smallest_elements(arrays: list[list[int]], k: int) -> list[int]:
    pass

Input Format

  • arrays: A list of lists of integers, where each arrays[i] is sorted in non-decreasing order.
  • k: A non-negative integer representing the maximum number of smallest elements to return.

Output Format

  • Return a list of integers containing the k smallest elements in non-decreasing order.

Constraints

  • 0 <= len(arrays) <= 10^4
  • 0 <= len(arrays[i]) <= 10^5
  • 0 <= k <= 10^5
  • -10^9 <= arrays[i][j] <= 10^9
  • Each arrays[i] is sorted in non-decreasing order.

Examples

Example 1

Input: arrays = [[1, 4, 7], [2, 5, 8], [3, 6, 9]], k = 5
Output: [1, 2, 3, 4, 5]
Explanation: Merging the arrays yields [1, 2, 3, 4, 5, 6, 7, 8, 9]. The 5 smallest elements are [1, 2, 3, 4, 5].

Example 2

Input: arrays = [[], [2, 2], [1, 3]], k = 4
Output: [1, 2, 2, 3]
Explanation: Empty lists are skipped. Duplicate values like 2 are preserved.

Example 3

Input: arrays = [[5, 10], [3]], k = 10
Output: [3, 5, 10]
Explanation: There are only 3 elements total across all arrays, which is fewer than k = 10, so all elements are returned.

Edge Cases

  • arrays is empty or only contains empty lists: return [].
  • k == 0: return [].
  • k is larger than the total number of elements across all arrays: return all elements sorted.

Hints

Hint 1

Since each array is already sorted, you only ever need to compare the current smallest unprocessed element from each array at any given step.

Hint 2

A min-heap storing tuples of `(value, array_index, element_index)` allows you to find and extract the current minimum across all arrays in `O(log m)` time.

Hint 3

Remember to initialize the heap with only non-empty arrays, and advance the index within that specific array each time an element is extracted.

Follow-up questions

  • How would you adapt your algorithm if the input arrays were iterators or infinite data streams that cannot be indexed directly?
  • If k is extremely small compared to m (e.g., k = 5 and m = 1,000,000), how could you optimize the initial heap construction?
  • How would your approach change if you were required to return only unique elements among the k smallest values?