3505. Minimum Operations to Make Elements Within K Subarrays Equal

/**
 * Problem: LeetCode 3505 - Minimum Operations to Make Elements Within K Subarrays Equal
 * Difficulty: Hard
 * Language: Java
 * Author: CodeByTushu
 */
import java.util.*;

class Solution {
    int m;
    long[] bc, bs;

    private void upd(int i, int c, long v) {
        for (; i <= m; i += i & -i) {
            bc[i] += c;
            bs[i] += v;
        }
    }

    public long minOperations(int[] nums, int x, int k) {
        int n = nums.length;

        // Coordinate compression
        int[] sorted = nums.clone();
        Arrays.sort(sorted);
        int u = 0;
        for (int i = 0; i < n; i++) {
            if (i == 0 || sorted[i] != sorted[i - 1]) sorted[u++] = sorted[i];
        }
        int[] vals = Arrays.copyOf(sorted, u);
        m = u;
        bc = new long[m + 1];
        bs = new long[m + 1];

        int[] idx = new int[n];
        for (int i = 0; i < n; i++) {
            idx[i] = Arrays.binarySearch(vals, nums[i]) + 1; // 1-indexed
        }

        int LOG = Integer.highestOneBit(m);
        long[] cost = new long[n - x + 1];
        long total = 0;

        for (int i = 0; i < n; i++) {
            upd(idx[i], 1, nums[i]);
            total += nums[i];

            if (i >= x) {
                upd(idx[i - x], -1, -nums[i - x]);
                total -= nums[i - x];
            }

            if (i >= x - 1) {
                int kth = x / 2 + 1;
                int pos = 0;
                long accSum = 0;
                int rem = kth;
                for (int pw = LOG; pw > 0; pw >>= 1) {
                    int np = pos + pw;
                    if (np <= m && bc[np] < rem) {
                        pos = np;
                        rem -= (int) bc[np];
                        accSum += bs[np];
                    }
                }
                long median = vals[pos]; // pos+1 in 1-indexed
                long sumLow = accSum + (long) rem * median;
                long sumUp = total - sumLow;

                cost[i - x + 1] = median * kth - sumLow + sumUp - median * (x - kth);
            }
        }

        long INF = Long.MAX_VALUE / 4;
        long[] prev = new long[n + 1]; // j = 0: cost 0
        for (int j = 1; j <= k; j++) {
            long[] cur = new long[n + 1];
            cur[0] = INF;
            for (int i = 1; i <= n; i++) {
                cur[i] = cur[i - 1];
                if (i >= x && prev[i - x] < INF) {
                    cur[i] = Math.min(cur[i], prev[i - x] + cost[i - x]);
                }
            }
            prev = cur;
        }

        return prev[n];
    }
}

"""
Problem: LeetCode 3505 - Minimum Operations to Make Elements Within K Subarrays Equal
Difficulty: Hard
Language: Python (Python 3)
Author: CodeByTushu
"""

from typing import List

class Solution:
    def minOperations(self, nums: List[int], x: int, k: int) -> int:
        n = len(nums)
        
        # Coordinate compression
        vals = sorted(list(set(nums)))
        m = len(vals)
        val_to_idx = {v: i + 1 for i, v in enumerate(vals)}
        idx = [val_to_idx[num] for num in nums]
        
        bc = [0] * (m + 1)
        bs = [0] * (m + 1)
        
        def upd(i: int, c: int, v: int):
            while i <= m:
                bc[i] += c
                bs[i] += v
                i += i & -i
                
        LOG = 1
        while (LOG << 1) <= m:
            LOG <<= 1
            
        cost = [0] * (n - x + 1)
        total = 0
        kth = x // 2 + 1
        
        for i in range(n):
            upd(idx[i], 1, nums[i])
            total += nums[i]
            
            if i >= x:
                upd(idx[i - x], -1, -nums[i - x])
                total -= nums[i - x]
                
            if i >= x - 1:
                pos = 0
                acc_sum = 0
                rem = kth
                pw = LOG
                while pw > 0:
                    np = pos + pw
                    if np <= m and bc[np] < rem:
                        pos = np
                        rem -= bc[np]
                        acc_sum += bs[np]
                    pw >>= 1
                    
                median = vals[pos]
                sum_low = acc_sum + rem * median
                sum_up = total - sum_low
                
                cost[i - x + 1] = median * kth - sum_low + sum_up - median * (x - kth)
                
        INF = float('inf')
        prev = [0] * (n + 1)
        for j in range(1, k + 1):
            cur = [INF] * (n + 1)
            for i in range(1, n + 1):
                cur[i] = cur[i - 1]
                if i >= x and prev[i - x] < INF:
                    cur[i] = min(cur[i], prev[i - x] + cost[i - x])
            prev = cur
            
        return prev[n]

/**
 * Problem: LeetCode 3505 - Minimum Operations to Make Elements Within K Subarrays Equal
 * Difficulty: Hard
 * Language: C++
 * Author: CodeByTushu
 */

#include <vector>
#include <algorithm>
#include <climits>

using namespace std;

class Solution {
    int m;
    vector<long long> bc, bs;

    void upd(int i, int c, long long v) {
        for (; i <= m; i += i & -i) {
            bc[i] += c;
            bs[i] += v;
        }
    }

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

        // Coordinate compression
        vector<int> vals = nums;
        sort(vals.begin(), vals.end());
        vals.erase(unique(vals.begin(), vals.end()), vals.end());
        m = vals.size();

        bc.assign(m + 1, 0);
        bs.assign(m + 1, 0);

        vector<int> idx(n);
        for (int i = 0; i < n; i++) {
            idx[i] = lower_bound(vals.begin(), vals.end(), nums[i]) - vals.begin() + 1;
        }

        int LOG = 1;
        while ((LOG << 1) <= m) LOG <<= 1;

        vector<long long> cost(n - x + 1, 0);
        long long total = 0;
        int kth = x / 2 + 1;

        for (int i = 0; i < n; i++) {
            upd(idx[i], 1, nums[i]);
            total += nums[i];

            if (i >= x) {
                upd(idx[i - x], -1, -nums[i - x]);
                total -= nums[i - x];
            }

            if (i >= x - 1) {
                int pos = 0;
                long long accSum = 0;
                int rem = kth;
                for (int pw = LOG; pw > 0; pw >>= 1) {
                    int np = pos + pw;
                    if (np <= m && bc[np] < rem) {
                        pos = np;
                        rem -= (int)bc[np];
                        accSum += bs[np];
                    }
                }
                long long median = vals[pos];
                long long sumLow = accSum + (long long)rem * median;
                long long sumUp = total - sumLow;

                cost[i - x + 1] = median * kth - sumLow + sumUp - median * (x - kth);
            }
        }

        long long INF = LLONG_MAX / 4;
        vector<long long> prev(n + 1, 0);
        for (int j = 1; j <= k; j++) {
            vector<long long> cur(n + 1, INF);
            for (int i = 1; i <= n; i++) {
                cur[i] = cur[i - 1];
                if (i >= x && prev[i - x] < INF) {
                    cur[i] = min(cur[i], prev[i - x] + cost[i - x]);
                }
            }
            prev = cur;
        }

        return prev[n];
    }
};

/**
 * Problem: LeetCode 3505 - Minimum Operations to Make Elements Within K Subarrays Equal
 * Difficulty: Hard
 * Language: JavaScript
 * Author: CodeByTushu
 */

/**
 * @param {number[]} nums
 * @param {number} x
 * @param {number} k
 * @return {number}
 */
var minOperations = function(nums, x, k) {
    const n = nums.length;

    // Coordinate compression
    const sorted = [...new Set(nums)].sort((a, b) => a - b);
    const m = sorted.length;
    const valMap = new Map();
    for (let i = 0; i < m; i++) valMap.set(sorted[i], i + 1);

    const idx = new Int32Array(n);
    for (let i = 0; i < n; i++) idx[i] = valMap.get(nums[i]);

    const bc = new Float64Array(m + 1);
    const bs = new Float64Array(m + 1);

    function upd(i, c, v) {
        for (; i <= m; i += i & -i) {
            bc[i] += c;
            bs[i] += v;
        }
    }

    let LOG = 1;
    while ((LOG << 1) <= m) LOG <<= 1;

    const cost = new Float64Array(n - x + 1);
    let total = 0;
    const kth = Math.floor(x / 2) + 1;

    for (let i = 0; i < n; i++) {
        upd(idx[i], 1, nums[i]);
        total += nums[i];

        if (i >= x) {
            upd(idx[i - x], -1, -nums[i - x]);
            total -= nums[i - x];
        }

        if (i >= x - 1) {
            let pos = 0;
            let accSum = 0;
            let rem = kth;
            for (let pw = LOG; pw > 0; pw >>= 1) {
                const np = pos + pw;
                if (np <= m && bc[np] < rem) {
                    pos = np;
                    rem -= bc[np];
                    accSum += bs[np];
                }
            }

            const median = sorted[pos];
            const sumLow = accSum + rem * median;
            const sumUp = total - sumLow;

            cost[i - x + 1] = median * kth - sumLow + sumUp - median * (x - kth);
        }
    }

    const INF = Number.MAX_SAFE_INTEGER / 4;
    let prev = new Float64Array(n + 1);
    for (let j = 1; j <= k; j++) {
        const cur = new Float64Array(n + 1);
        cur[0] = INF;
        for (let i = 1; i <= n; i++) {
            cur[i] = cur[i - 1];
            if (i >= x && prev[i - x] < INF) {
                cur[i] = Math.min(cur[i], prev[i - x] + cost[i - x]);
            }
        }
        prev = cur;
    }

    return prev[n];
};