Maximum Subsequence Score — Medium Problem & Solution

Two arrays of the same length are given. Choose a subsequence of exactly k indices.

  • Difficulty: Medium
  • Topics: Arrays, Greedy, Sorting, Heap
  • Asked at: Amazon, Google, Flipkart
  • Time limit: 2 s
  • Memory limit: 256 MB
  • Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

Two arrays of the same length are given. Choose a subsequence of exactly k indices. Its score is the sum of the chosen nums1 values multiplied by the minimum of the chosen nums2 values.

Return the maximum score.

Example 1

Input: nums1 = [1,3,3,2], nums2 = [2,1,3,4], k = 3
Output: 12
Explanation: Indices 0, 2 and 3: `(1 + 3 + 2) × min(2, 3, 4) = 6 × 2`.

Example 2

Input: nums1 = [4,2,3,1,1], nums2 = [7,5,10,9,6], k = 1
Output: 30
Explanation: Index 2 alone: `3 × 10`.

Example 3

Input: nums1 = [1,1], nums2 = [5,5], k = 2
Output: 10

Constraints

  • n == nums1.length == nums2.length
  • 1 <= n <= 1000
  • 0 <= nums1[i], nums2[j] <= 1000
  • 1 <= k <= n

How to solve Maximum Subsequence Score

Sort by nums2 descending. Sweeping in that order, the element just added is the smallest nums2 among those considered, so it fixes the multiplier; the best sum is then simply the k largest nums1 values seen, which a size-k min-heap maintains.

Approach

  1. Sort the indices by nums2, largest first.
  2. Push each nums1 value onto a min-heap, tracking the running sum.
  3. If the heap exceeds k, pop the smallest and subtract it.
  4. Once the heap holds exactly k, score sum × nums2[current] and keep the best.

Why it works

Sorting descending is what lets the multiplier be read off for free: every element already in the heap has a nums2 at least as large as the current one, so the current one is the minimum. Dropping the smallest nums1 when the heap overflows is safe because it can never be part of the best sum for this or any later multiplier.

Complexity

  • Time — O(n log n)
  • Space — O(n)

Pitfalls

  • The subsequence must hold exactly k indices, so scores are only taken once the heap is full.
  • Sorting ascending by nums2 breaks the "current element is the minimum" invariant.
  • The products reach 10^6 × 10^3, so a 32-bit int is enough here but only just.

Reference solution

Python

from typing import List
import heapq

def maxScore(nums1: List[int], nums2: List[int], k: int) -> int:
    order = sorted(range(len(nums1)), key=lambda i: -nums2[i])
    heap = []
    total = 0
    best = 0
    for i in order:
        heapq.heappush(heap, nums1[i])
        total += nums1[i]
        if len(heap) > k:
            total -= heapq.heappop(heap)
        if len(heap) == k:
            best = max(best, total * nums2[i])
    return best

JavaScript

var maxScore = function(nums1, nums2, k) {
    var n = nums1.length, i;
    var order = [];
    for (i = 0; i < n; i++) order.push(i);
    order.sort(function(a, b) { return nums2[b] - nums2[a]; });
    var heap = [];
    var push = function(v) {
        heap.push(v);
        var j = heap.length - 1;
        while (j > 0) {
            var p = (j - 1) >> 1;
            if (heap[p] <= heap[j]) break;
            var t = heap[p]; heap[p] = heap[j]; heap[j] = t;
            j = p;
        }
    };
    var pop = function() {
        var top = heap[0];
        var last = heap.pop();
        if (heap.length > 0) {
            heap[0] = last;
            var j = 0;
            for (;;) {
                var l = 2 * j + 1, r = l + 1, s = j;
                if (l < heap.length && heap[l] < heap[s]) s = l;
                if (r < heap.length && heap[r] < heap[s]) s = r;
                if (s === j) break;
                var t = heap[s]; heap[s] = heap[j]; heap[j] = t;
                j = s;
            }
        }
        return top;
    };
    var sum = 0, best = 0;
    for (var t2 = 0; t2 < order.length; t2++) {
        var idx = order[t2];
        push(nums1[idx]);
        sum += nums1[idx];
        if (heap.length > k) sum -= pop();
        if (heap.length === k) {
            var score = sum * nums2[idx];
            if (score > best) best = score;
        }
    }
    return best;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 667 arrays problems · the whole catalogue