I'm trying to improve my Python, so I'm trying to do coding challenges and getting them reviewed. I would really appreciate if someone can point out how I can make my code cleaner/more pythonic. My min_heap with a 3-value tuple seems a bit messy, but I couldn't think of a better way to do it.
Problem:
https://leetcode.com/problems/find-k-pairs-with-smallest-sums/description/
You are given two integer arrays
nums1
andnums2
sorted in ascending order and an integerk
.Define a pair
(u,v)
which consists of one element from the first array and one element from the second array.Find the
k
pairs(u1,v1),(u2,v2) ...(uk,vk)
with the smallest sums.
Example 1:
Given
nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Return:
[1,2],[1,4],[1,6]
The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
Solution:
from heapq import heappop, heappush
class Solution:
def kSmallestPairs(self, nums1, nums2, k):
"""
:type nums1: List[int]
:type nums2: List[int]
:type k: int
:rtype: List[List[int]]
"""
solution = []
min_heap = []
# if nums1 or nums2 is empty or k is 0, we're done
if not nums1 or not nums2 or not k:
return solution
# first initialize min_heap with all
# (nums1[0..n], nums2[0]) pairs with
# (sum of pair, (nums1[0..n], nums2[0]), index of nums2)
for value in nums1:
heappush(min_heap, (value+nums2[0], (value, nums2[0]), 0))
while k and min_heap:
curr_pair = heappop(min_heap)
solution.append(curr_pair[1])
nums2_idx = curr_pair[2]
k -= 1
# if we haven't exhausted all (nums1[curr], nums2[0..n])
# pairs, offer pair (nums1[curr], nums2[curr+1]) to min_heap
if curr_pair[2] == len(nums2)-1: continue
heappush(min_heap, (curr_pair[1][0]+nums2[nums2_idx+1], (curr_pair[1][0],
nums2[nums2_idx+1]), nums2_idx+1))
return solution