Skip to content

3962. Maximum Subarray Sum After at Most K Swaps

SourceWeekly Contest 506 Q4DifficultyHardRating2672

Description

You are given an integer array nums and an integer k.

You are allowed to perform at most k swap operations on the array.

In one swap operation, you may choose any two indices i and j and swap nums[i] and nums[j].

Return an integer denoting the maximum possible subarray sum after performing the swaps.

 

Example 1:

Input: nums = [1,-1,0,2], k = 1

Output: 3

Explanation:

  • We can swap on indices 1 and 3, resulting in the array [1, 2, 0, -1].
  • The subarray [1, 2] has a sum of 3, which is the maximum possible subarray sum after at most k = 1​​​​​​​ swap.

Example 2:

Input: nums = [4,3,2,4], k = 2

Output: 13

Explanation:

The maximum possible subarray sum after at most k = 2 swaps is the sum of the entire array, which is 13.

Example 3:

Input: nums = [-1,-2], k = 0

Output: -1

Explanation:

  • k = 0 swaps are allowed.
  • The possible subarrays are [-1], [-2], and [-1, -2], with sums -1, -2, and -3 respectively.
  • Among these sums, the maximum is -1.

 

Constraints:

  • 1 <= nums.length <= 1500
  • -105 <= nums[i] <= 105
  • 0 <= k <= nums.length

Solutions

Solution 1: Enumerate Subarrays + Binary Indexed Tree

Thinking

A swap helps the chosen subarray only when it brings in a larger value from outside. With \(n \le 1500\) every segment can be enumerated; sorting the inside and the outside of each one would add another \(O(n \log n)\) and push the total to \(O(n^3 \log n)\).

Once the inside is ordered ascending and the outside descending, the gains are monotone. If the \(t\)-th pair still has the smaller value inside, the first \(t\) pairs are all worth swapping; once a pair does not increase the sum, later pairs do not either. How many swaps to take can be found by binary search.

After coordinate compression, two Fenwick trees store the counts and sums inside and outside the segment, so the smallest or largest values can be read by rank. With the left end fixed, extending the right end moves the current value from the outside tree into the inside tree.

Deduplicate and sort the values in \(\textit{nums}\), and use that list as the compressed value domain. Two Fenwick trees keep, for the inside and the outside, how many times each value occurs and what those occurrences sum to. Either tree can report the \(k\)-th smallest value, or the sum of the smallest or largest several values, in \(O(\log n)\).

Enumerate the left end \(l\). Every element starts outside the segment. The right end \(r\) runs from \(l\) to \(n - 1\): move \(\textit{nums}[r]\) from the outside tree into the inside tree and add it to the segment sum \(s\).

Let \(c\) be the number of elements inside. At most \(t = \min(k, c, n - c)\) swaps are possible. Binary search the largest \(mid\) in \([1, t]\) such that the \(mid\)-th smallest inside value is strictly smaller than the \(mid\)-th largest outside value, and call it \(best\). If \(best > 0\), a candidate is \(s\) plus the sum of the \(best\) largest outside values, minus the sum of the \(best\) smallest inside values. The answer is the maximum candidate over all segments.

The time complexity is \(O(n^2 \log^2 n)\) and the space complexity is \(O(n)\), where \(n\) is the length of \(\textit{nums}\).

1


  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
class Fenwick {
    int n;
    int[] count;
    long[] sum;

    Fenwick(int n) {
        this.n = n;
        count = new int[n + 1];
        sum = new long[n + 1];
    }

    void add(int idx, int cnt, long val) {
        ++idx;
        while (idx <= n) {
            count[idx] += cnt;
            sum[idx] += val;
            idx += idx & -idx;
        }
    }

    int prefixCount(int idx) {
        int res = 0;
        while (idx > 0) {
            res += count[idx];
            idx -= idx & -idx;
        }
        return res;
    }

    long prefixSum(int idx) {
        long res = 0;
        while (idx > 0) {
            res += sum[idx];
            idx -= idx & -idx;
        }
        return res;
    }

    int kth(int k) {
        int idx = 0;
        for (int bit = Integer.highestOneBit(n); bit > 0; bit >>= 1) {
            int next = idx + bit;
            if (next <= n && count[next] < k) {
                idx = next;
                k -= count[next];
            }
        }
        return idx;
    }

    long sumSmallest(int k, int[] values) {
        if (k <= 0) {
            return 0;
        }
        int pos = kth(k);
        int before = prefixCount(pos);
        return prefixSum(pos) + (long) (k - before) * values[pos];
    }

    long sumLargest(int k, int[] values) {
        int total = prefixCount(n);
        if (k <= 0) {
            return 0;
        }
        if (k >= total) {
            return prefixSum(n);
        }
        return prefixSum(n) - sumSmallest(total - k, values);
    }
}

class Solution {
    public long maxSum(int[] nums, int k) {
        int n = nums.length;
        int[] sorted = nums.clone();
        Arrays.sort(sorted);
        int m = 0;
        for (int i = 0; i < n; ++i) {
            if (i == 0 || sorted[i] != sorted[i - 1]) {
                sorted[m++] = sorted[i];
            }
        }
        int[] values = Arrays.copyOf(sorted, m);
        int[] idx = new int[n];
        for (int i = 0; i < n; ++i) {
            idx[i] = Arrays.binarySearch(values, nums[i]);
        }

        long ans = Long.MIN_VALUE;
        for (int l = 0; l < n; ++l) {
            Fenwick inside = new Fenwick(m);
            Fenwick outside = new Fenwick(m);
            for (int i = 0; i < n; ++i) {
                outside.add(idx[i], 1, nums[i]);
            }
            long window = 0;
            for (int r = l; r < n; ++r) {
                outside.add(idx[r], -1, -nums[r]);
                inside.add(idx[r], 1, nums[r]);
                window += nums[r];
                int inCnt = r - l + 1;
                int outCnt = n - inCnt;
                int limit = Math.min(k, Math.min(inCnt, outCnt));
                if (limit == 0) {
                    ans = Math.max(ans, window);
                    continue;
                }
                int lo = 1;
                int hi = limit;
                int best = 0;
                while (lo <= hi) {
                    int mid = (lo + hi) >>> 1;
                    int small = inside.kth(mid);
                    int large = outside.kth(outCnt - mid + 1);
                    if (values[small] < values[large]) {
                        best = mid;
                        lo = mid + 1;
                    } else {
                        hi = mid - 1;
                    }
                }
                if (best == 0) {
                    ans = Math.max(ans, window);
                } else {
                    long cand = window + outside.sumLargest(best, values)
                        - inside.sumSmallest(best, values);
                    ans = Math.max(ans, cand);
                }
            }
        }
        return ans;
    }
}

1


1

Comments