Skip to content

Instantly share code, notes, and snippets.

@chrisbodhi
Created December 16, 2023 14:12
Show Gist options
  • Save chrisbodhi/4c787309a12c84696ad820953c143674 to your computer and use it in GitHub Desktop.
Save chrisbodhi/4c787309a12c84696ad820953c143674 to your computer and use it in GitHub Desktop.
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment