find_k_mins.py

import heapq

def find_k_mins(matrix, k):
  ptrs = [0 for _ in matrix]
  curr_min = float("inf")
  min_ind = -1
  heap = []
  heapq.heapify(heap)

  while len(heap) < k:
    for i in range(len(ptrs)):
      ptr = ptrs[i]
      val = matrix[i][ptr]
      if val <= curr_min:
        curr_min = val
        min_ind = i
    ptrs[min_ind] += 1
    heapq.heappush(heap, curr_min)
    curr_min = float("inf")
    min_ind = -1

  return heap



def runner():
  scenarios = [
    {
      "matrix": [],
      "k": 0,
      "expected": []
    },
    {
      "matrix": [
        [1, 2, 4],
        [2, 3, 4],
        [4, 5, 7]
      ],
      "k": 4,
      "expected": [1, 2, 2, 3]
    },
    {
      "matrix": [
        [9, 12, 14, 17, 22],
        [4, 8, 11, 13, 71],
        [2, 6, 7, 8, 9]
      ],
      "k": 6,
      "expected": [2, 4, 6, 7, 8, 8]
    }
  ]

  for scenario in scenarios:
    l = find_k_mins(scenario['matrix'], scenario['k'])
    if l != scenario['expected']:
      print(f"Got {l}, expected {scenario['expected']} with a k of {scenario['k']}")
    else:
      print("+1")

runner()