Skip to content

3505. Minimum Operations to Make Elements Within K Subarrays Equal 👍

  • Time: $O(nk)$
  • Space: $O(nk)$
 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
class Solution {
 public:
  long long minOperations(vector<int>& nums, int x, int k) {
    // minOps[i] := the minimum number of operations needed to make
    // nums[i..i + x - 1] equal to the median
    const vector<long> minOps = getMinOps(nums, x);
    vector<vector<long>> mem(nums.size() + 1, vector<long>(k + 1, -1));
    return minOperations(nums, x, 0, k, minOps, mem);
  }

 private:
  static constexpr long kInf = LONG_MAX / 2;

  // Returns the minimum operations needed to have at least k non-overlapping
  // subarrays of size x in nums[i..n - 1].
  long minOperations(const vector<int>& nums, int x, int i, int k,
                     const vector<long>& minOps, vector<vector<long>>& mem) {
    if (k == 0)
      return 0;
    if (i == nums.size())
      return kInf;
    if (mem[i][k] != -1)
      return mem[i][k];
    const long skip = minOperations(nums, x, i + 1, k, minOps, mem);
    const long pick =
        i + x <= nums.size()
            ? minOps[i] + minOperations(nums, x, i + x, k - 1, minOps, mem)
            : kInf;
    return mem[i][k] = min(skip, pick);
  }

  // Returns the minimum operations needed to make all elements in the window of
  // size x equal to the median.
  vector<long> getMinOps(const vector<int>& nums, int x) {
    vector<long> minOps;
    multiset<int> lower;
    multiset<int> upper;
    long lowerSum = 0;
    long upperSum = 0;
    for (int i = 0; i < nums.size(); ++i) {
      if (lower.empty() || nums[i] <= *lower.rbegin()) {
        lower.insert(nums[i]);
        lowerSum += nums[i];
      } else {
        upper.insert(nums[i]);
        upperSum += nums[i];
      }
      if (i >= x) {
        const int outNum = nums[i - x];
        if (const auto it = lower.find(outNum); it != lower.cend()) {
          lower.erase(it);
          lowerSum -= outNum;
        } else {
          upper.erase(upper.find(outNum));
          upperSum -= outNum;
        }
      }
      // Balance the two multisets s.t.
      // |lower| >= |upper| and |lower| - |upper| <= 1.
      if (lower.size() < upper.size()) {
        const int val = *upper.begin();
        upper.erase(upper.begin());
        lower.insert(val);
        upperSum -= val;
        lowerSum += val;
      } else if (lower.size() - upper.size() > 1) {
        const int val = *lower.rbegin();
        lower.erase(prev(lower.end()));
        upper.insert(val);
        lowerSum -= val;
        upperSum += val;
      }
      // Calculate operations needed to make all elements in the window equal
      // to the median.
      if (i >= x - 1) {
        const int median = *lower.rbegin();
        const long ops = (median * lower.size() - lowerSum) +
                         (upperSum - median * upper.size());
        minOps.push_back(ops);
      }
    }
    return minOps;
  }
};

 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
class MyMap {
  public TreeMap<Integer, Integer> map = new TreeMap<>();
  public int size = 0;
  public long sum = 0;
}

class Solution {
  public long minOperations(int[] nums, int x, int k) {
    // minOps[i] := the minimum number of operations needed to make
    // nums[i..i + x - 1] equal to the median
    List<Long> minOps = getMinOps(nums, x);
    Long[][] mem = new Long[nums.length + 1][k + 1];
    return minOperations(nums, x, 0, k, minOps, mem);
  }

  private static final long INF = Long.MAX_VALUE / 2;

  // Returns the minimum operations needed to have at least k non-overlapping
  // subarrays of size x in nums[i..n - 1].
  private long minOperations(int[] nums, int x, int i, int k, List<Long> minOps, Long[][] mem) {
    if (k == 0)
      return 0;
    if (i == nums.length)
      return INF;
    if (mem[i][k] != null)
      return mem[i][k];
    final long skip = minOperations(nums, x, i + 1, k, minOps, mem);
    final long pick = i + x <= nums.length
                          ? minOps.get(i) + minOperations(nums, x, i + x, k - 1, minOps, mem)
                          : INF;
    return mem[i][k] = Math.min(skip, pick);
  }

  // Returns the minimum operations needed to make all elements in the window of
  // size x equal to the median.
  private List<Long> getMinOps(int[] nums, int x) {
    List<Long> minOps = new ArrayList<>();
    MyMap lower = new MyMap();
    MyMap upper = new MyMap();
    for (int i = 0; i < nums.length; ++i) {
      if (lower.map.isEmpty() || nums[i] <= lower.map.lastKey()) {
        lower.map.merge(nums[i], 1, Integer::sum);
        lower.sum += nums[i];
        ++lower.size;
      } else {
        upper.map.merge(nums[i], 1, Integer::sum);
        upper.sum += nums[i];
        ++upper.size;
      }
      if (i >= x) {
        final int outNum = nums[i - x];
        if (lower.map.containsKey(outNum)) {
          lower.map.merge(outNum, -1, Integer::sum);
          if (lower.map.get(outNum) == 0)
            lower.map.remove(outNum);
          lower.sum -= outNum;
          --lower.size;
        } else {
          upper.map.merge(outNum, -1, Integer::sum);
          if (upper.map.get(outNum) == 0)
            upper.map.remove(outNum);
          upper.sum -= outNum;
          --upper.size;
        }
      }
      // Balance the two maps s.t.
      // |lower| >= |upper| and |lower| - |upper| <= 1.
      if (lower.size < upper.size) {
        final int val = upper.map.firstKey();
        upper.map.merge(val, -1, Integer::sum);
        if (upper.map.get(val) == 0)
          upper.map.remove(val);
        lower.map.merge(val, 1, Integer::sum);
        upper.sum -= val;
        lower.sum += val;
        --upper.size;
        ++lower.size;
      } else if (lower.size - upper.size > 1) {
        final int val = lower.map.lastKey();
        lower.map.merge(val, -1, Integer::sum);
        if (lower.map.get(val) == 0)
          lower.map.remove(val);
        upper.map.merge(val, 1, Integer::sum);
        lower.sum -= val;
        upper.sum += val;
        --lower.size;
        ++upper.size;
      }
      // Calculate operations needed to make all elements in the window equal
      // to the median.
      if (i >= x - 1) {
        final int median = lower.map.lastKey();
        final long ops = (median * lower.size - lower.sum) + (upper.sum - median * upper.size);
        minOps.add(ops);
      }
    }
    return minOps;
  }
}