← leetcode

Minimum Operations to Make Elements Within K Subarrays Equal

· Hard

LeetCode problem 3505

intuitively speaking, the logic of this question does not seem too bad! it already says that we have to find

the minimum number of operations needed to have at least k non-overlapping subarrays of size exactly x in nums, where all elements within each subarray are equal

using information from a previous question, we understand that for any given subarray, the value that takes the fewest operations to make all elements equal to is simply the median value inside the subarray, so the minimum number of operations is the sum of each element's distance to the median.

therefore, if we precalculate the minimum operations for every subarray of size x, we can use dynamic programming to calculate the minimum total number of operations.

calculating each subarray from scratch is too slow though, since there are n - x + 1 subarrays of x elements each. instead, we can slide the window one element at a time, keeping its elements split into two sorted halves along with each half's sum:

  • the median is the largest element of the smaller half
  • the operations needed are median * (size of smaller half) - (sum of smaller half) plus (sum of larger half) - median * (size of larger half)
  • when the window moves, we remove the element that leaves, add the element that arrives, and rebalance the two halves

this takes O(log(x)) per subarray instead of O(x).

for the dynamic programming, dp[i][j] is the minimum operations to make j subarrays using only the first i elements: either no subarray ends at element i, giving dp[i-1][j], or one does, giving dp[i-x][j-1] plus that subarray's operations. we only ever need exactly k subarrays, since an extra one can only add more operations.

code link to code

template <typename Compare>
class LazyHeap {
public:
    void push(int x) {
        pq_.push(x);
        ++size_;
    }

    void erase(int x) {
        ++to_delete_[x];
        --size_;
        prune();
    }

    int top() {
        prune();
        return pq_.top();
    }

    void pop() {
        prune();
        pq_.pop();
        --size_;
    }

    int size() {
        return size_;
    }

private:
    void prune() {
        while (!pq_.empty()) {
            auto it = to_delete_.find(pq_.top());
            if (it == to_delete_.end()) return;
            if (--it->second == 0) {
                to_delete_.erase(it);
            }
            pq_.pop();
        }
    }

    priority_queue<int, vector<int>, Compare> pq_;
    unordered_map<int, int> to_delete_;
    int size_{};
};

class SlidingWindow {
public:
    void add(int n) {
        if (left_.size() == 0 || n <= left_.top()) {
            left_.push(n);
            left_total_ += n;
        } else {
            right_.push(n);
            right_total_ += n;
        }
        resize();
    }

    void remove(int n) {
        if (n <= left_.top()) {
            left_.erase(n);
            left_total_ -= n;
        } else {
            right_.erase(n);
            right_total_ -= n;
        }
        resize();
    }

    long long operations() {
        if (left_.size() == 0) return 0;
        long long median = left_.top();

        return (left_.size() * median - left_total_) 
            + (right_total_ - right_.size() * median);
    } 
private:
    void resize() {
        if (left_.size() > right_.size() + 1) {
            int n = left_.top();
            left_.pop();
            left_total_ -= n;
            right_.push(n);
            right_total_ += n;
        } else if (left_.size() < right_.size()) {
            int n = right_.top();
            right_.pop();
            right_total_ -= n;
            left_.push(n);
            left_total_ += n;
        }
    }

    LazyHeap<less<int>> left_;
    LazyHeap<greater<int>> right_;
    long long left_total_{}, right_total_{};
};

class Solution {
public:
    long long minOperations(vector<int>& nums, int x, int k) {
        int n {static_cast<int>(nums.size())};

        SlidingWindow sw;
        vector<long long> cost (n - x + 1);

        for (int i = 0; i < x; ++i) sw.add(nums[i]);
        cost[0] = sw.operations();

        for (int i = x; i < n; ++i) {
            sw.add(nums[i]);
            sw.remove(nums[i - x]);
            cost[i - x + 1] = sw.operations();
        }

        vector<long long> prev (n + 1, 0);
        for (int j = 1; j <= k; ++j) {
            vector<long long> cur (n + 1, LLONG_MAX / 2);
            for (int i = x; i <= n; ++i) {
                long long skip = cur[i-1];
                long long take = prev[i-x] + cost[i-x];
                cur[i] = min(skip, take);
            }
            prev = move(cur);
        }

        return prev[n];
    }
};

complexity link to complexity

  • Time: O(n log(n) + nk), where n denotes nums.size()
  • Space: O(n)

time taken link to time taken

  • DNF