Maximum Total Subarray Value II

hard max-heap sparse table greedy

Problem

Given an array nums of length n and an integer k, choose exactly k distinct subarrays nums[l..r] (overlap allowed, but no repeated (l, r)). The value of a subarray is max(nums[l..r]) − min(nums[l..r]). Return the maximum possible total value, i.e. the sum of the values of the chosen subarrays.

Inputnums = [1, 3, 2], k = 2

Output4

Pick [0..1] = [1,3] (value 3−1 = 2) and [0..2] = [1,3,2] (value 3−1 = 2). Total = 2 + 2 = 4.

nums (comma-separated)

import heapq

def maxTotalValue(nums, k):
    n = len(nums)
    # Sparse tables for O(1) range max / range min.
    LOG = max(1, n.bit_length())
    mx = [nums[:]]; mn = [nums[:]]
    j = 1
    while (1 << j) <= n:
        pmx, pmn = mx[-1], mn[-1]; half = 1 << (j - 1)
        cmx, cmn = [], []
        for i in range(n - (1 << j) + 1):
            cmx.append(max(pmx[i], pmx[i + half]))
            cmn.append(min(pmn[i], pmn[i + half]))
        mx.append(cmx); mn.append(cmn); j += 1

    def value(l, r):                       # max(nums[l..r]) - min(...)
        s = (r - l + 1).bit_length() - 1
        hi = max(mx[s][l], mx[s][r - (1 << s) + 1])
        lo = min(mn[s][l], mn[s][r - (1 << s) + 1])
        return hi - lo

    # For fixed l, value(l, r) is largest at r = n-1, so seed the heap there.
    heap = [(-value(l, n - 1), l, n - 1) for l in range(n)]
    heapq.heapify(heap)

    total = 0
    for _ in range(k):                     # take the k biggest values
        v, l, r = heapq.heappop(heap)
        total += -v
        if r > l:                          # next best for this l: shrink r
            heapq.heappush(heap, (-value(l, r - 1), l, r - 1))
    return total

function maxTotalValue(nums, k) {
  const n = nums.length;
  // Sparse tables for O(1) range max / range min.
  const mx = [nums.slice()], mn = [nums.slice()];
  for (let j = 1; (1 << j) <= n; j++) {
    const pmx = mx[j - 1], pmn = mn[j - 1], half = 1 << (j - 1);
    const cmx = [], cmn = [];
    for (let i = 0; i + (1 << j) <= n; i++) {
      cmx.push(Math.max(pmx[i], pmx[i + half]));
      cmn.push(Math.min(pmn[i], pmn[i + half]));
    }
    mx.push(cmx); mn.push(cmn);
  }
  const value = (l, r) => {               // max(nums[l..r]) - min(...)
    const s = 31 - Math.clz32(r - l + 1);
    const hi = Math.max(mx[s][l], mx[s][r - (1 << s) + 1]);
    const lo = Math.min(mn[s][l], mn[s][r - (1 << s) + 1]);
    return hi - lo;
  };
  // For fixed l, value(l, r) peaks at r = n-1; seed the max-heap there.
  const heap = new MaxHeap();
  for (let l = 0; l < n; l++) heap.push([value(l, n - 1), l, n - 1]);

  let total = 0n;
  for (let t = 0; t < k; t++) {           // take the k biggest values
    const [v, l, r] = heap.pop();
    total += BigInt(v);
    if (r > l) heap.push([value(l, r - 1), l, r - 1]);  // shrink r
  }
  return total;
}

long maxTotalValue(int[] nums, int k) {
    int n = nums.length, LOG = 32 - Integer.numberOfLeadingZeros(n);
    int[][] mx = new int[LOG][], mn = new int[LOG][];
    mx[0] = nums.clone(); mn[0] = nums.clone();
    for (int j = 1; (1 << j) <= n; j++) {              // build sparse tables
        int len = n - (1 << j) + 1, half = 1 << (j - 1);
        mx[j] = new int[len]; mn[j] = new int[len];
        for (int i = 0; i < len; i++) {
            mx[j][i] = Math.max(mx[j - 1][i], mx[j - 1][i + half]);
            mn[j][i] = Math.min(mn[j - 1][i], mn[j - 1][i + half]);
        }
    }
    // For fixed l, value(l, r) peaks at r = n-1; seed the max-heap there.
    PriorityQueue<long[]> heap = new PriorityQueue<>((a, b) -> Long.compare(b[0], a[0]));
    for (int l = 0; l < n; l++) heap.add(new long[]{value(mx, mn, l, n - 1), l, n - 1});

    long total = 0;
    for (int t = 0; t < k; t++) {                      // take the k biggest values
        long[] top = heap.poll();
        total += top[0];
        long l = top[1], r = top[2];
        if (r > l) heap.add(new long[]{value(mx, mn, (int) l, (int) r - 1), l, r - 1});
    }
    return total;
}

long value(int[][] mx, int[][] mn, int l, int r) {     // max(nums[l..r]) - min(...)
    int s = 31 - Integer.numberOfLeadingZeros(r - l + 1);
    return Math.max(mx[s][l], mx[s][r - (1 << s) + 1])
         - Math.min(mn[s][l], mn[s][r - (1 << s) + 1]);
}

long long maxTotalValue(vector<int>& nums, int k) {
    int n = nums.size(), LOG = 32 - __builtin_clz(n);
    vector<vector<int>> mx(LOG), mn(LOG);
    mx[0] = nums; mn[0] = nums;
    for (int j = 1; (1 << j) <= n; j++) {              // build sparse tables
        int len = n - (1 << j) + 1, half = 1 << (j - 1);
        mx[j].resize(len); mn[j].resize(len);
        for (int i = 0; i < len; i++) {
            mx[j][i] = max(mx[j - 1][i], mx[j - 1][i + half]);
            mn[j][i] = min(mn[j - 1][i], mn[j - 1][i + half]);
        }
    }
    auto value = [&](int l, int r) -> long long {      // max(nums[l..r]) - min(...)
        int s = 31 - __builtin_clz(r - l + 1);
        return max(mx[s][l], mx[s][r - (1 << s) + 1])
             - min(mn[s][l], mn[s][r - (1 << s) + 1]);
    };
    // For fixed l, value(l, r) peaks at r = n-1; seed the max-heap there.
    priority_queue<array<long long, 3>> heap;
    for (int l = 0; l < n; l++) heap.push({value(l, n - 1), l, n - 1});

    long long total = 0;
    for (int t = 0; t < k; t++) {                      // take the k biggest values
        auto [v, l, r] = heap.top(); heap.pop();
        total += v;
        if (r > l) heap.push({value(l, r - 1), l, r - 1});  // shrink r
    }
    return total;
}

Explanation

There are n(n+1)/2 subarrays — far too many to enumerate when n reaches 50 000. The trick is to select the k largest values without ever materialising every subarray.

Key monotonicity. Fix the left endpoint l. As the right endpoint r grows, the window nums[l..r] only gains elements, so the running maximum can only rise and the running minimum can only fall. Therefore value(l, r) = max − min is non-decreasing in r. For a fixed l, the single best value lives at r = n−1, and shrinking r gives the next-best, then the next, and so on.

Heap of candidates. Seed a max-heap with one entry per left endpoint: (value(l, n−1), l, n−1). Pop the global maximum, add it to the running total, and if that entry still has room to shrink (r > l) push its successor (value(l, r−1), l, r−1). Repeat k times. Because each l-chain is sorted and we always expose only its current frontier, the k pops are exactly the k largest values across all subarrays — a classic "merge k sorted lists" pattern.

Fast range queries. Each candidate needs max and min over an arbitrary window. We precompute two sparse tables so any value(l, r) is answered in O(1) by combining two overlapping power-of-two blocks. This keeps every heap operation cheap.

Why greedy is safe. Subarrays are distinct by (l, r), and the values we expose per l are strictly the sorted tail of that chain. Picking the top k from a collection of sorted lists, advancing only the list we drew from, provably yields the maximum sum of k elements. The answer can exceed 32 bits, so it is returned as a 64-bit long.

Time: O(n log n + k log n) Space: O(n log n)