less than 1 minute read

Problem Statement

leetcode problem link

My Solution [TLE]

class Solution:
    def maxScore(self, nums1: List[int], nums2: List[int], k: int) -> int:
        max_heap = []
        N = len(nums1)
        permutations = []
        def dfs(i, curr):
            if len(curr) == k:
                score = sum(nums1[i] for i in curr) * min(nums2[i] for i in curr)
                heapq.heappush(max_heap, -score)
                return
            for j in range(i, N):
                curr.append(j)
                dfs(j + 1, curr)
                curr.pop()

        dfs(0, [])

        return -max_heap[0]

Editorial

class Solution:
    def maxScore(self, nums1: List[int], nums2: List[int], k: int) -> int:
        # Sort pair (nums1[i], nums2[i]) by nums2[i] in decreasing order.
        pairs = [(a, b) for a, b in zip(nums1, nums2)]
        pairs.sort(key = lambda x: -x[1])

        # Use a min-heap to maintain the top k elements.
        top_k_heap = [x[0] for x in pairs[:k]]
        top_k_sum = sum(top_k_heap)
        heapq.heapify(top_k_heap)

        # The score of the first k pairs.
        answer = top_k_sum * pairs[k - 1][1]

        # Iterate over every nums2[i] as minimum from nums2.
        for i in range(k, len(nums1)):
            # Remove the smallest integer from the previous top k elements
            # then add nums1[i] to the top k elements.
            top_k_sum -= heapq.heappop(top_k_heap)
            top_k_sum += pairs[i][0]
            heapq.heappush(top_k_heap, pairs[i][0])

            # Update answer as the maximum score.
            answer = max(answer, top_k_sum * pairs[i][1])

        return answer