3013. Divide an Array Into Subarrays With Minimum Cost II

/**
 * Problem: LeetCode 3013 - Divide an Array Into Subarrays With Minimum Cost II
 * Difficulty: Hard
 * Language: Java
 * Author: CodeByTushu
 */
import java.util.Comparator;
import java.util.TreeSet;

class Solution {
    static class Pair {
        long val;
        int id;

        Pair(long val, int id) {
            this.val = val;
            this.id = id;
        }
    }

    private TreeSet<Pair> left;
    private TreeSet<Pair> right;
    private long sumL;
    private int m;

    public long minimumCost(int[] nums, int k, int dist) {
        int n = nums.length;
        this.m = k - 1;

        // Comparator ordering by value ascending, then by original index ascending
        Comparator<Pair> cmp = (a, b) -> {
            if (a.val != b.val) return Long.compare(a.val, b.val);
            return Integer.compare(a.id, b.id);
        };

        left = new TreeSet<>(cmp);
        right = new TreeSet<>(cmp);
        sumL = 0;

        // Initialize the first sliding window of size dist + 1: indices 1 to min(n - 1, 1 + dist)
        for (int i = 1; i <= Math.min(n - 1, 1 + dist); i++) {
            add(new Pair(nums[i], i), cmp);
        }

        long minSum = sumL;

        // Slide the window across nums[1 ... n - 1]
        for (int i = 2 + dist; i < n; i++) {
            add(new Pair(nums[i], i), cmp);
            remove(new Pair(nums[i - dist - 1], i - dist - 1));
            minSum = Math.min(minSum, sumL);
        }

        return (long) nums[0] + minSum;
    }

    private void add(Pair p, Comparator<Pair> cmp) {
        if (left.size() < m) {
            left.add(p);
            sumL += p.val;
        } else {
            if (cmp.compare(p, left.last()) < 0) {
                Pair maxL = left.pollLast();
                sumL -= maxL.val;
                right.add(maxL);

                left.add(p);
                sumL += p.val;
            } else {
                right.add(p);
            }
        }
    }

    private void remove(Pair p) {
        if (left.contains(p)) {
            left.remove(p);
            sumL -= p.val;
            if (!right.isEmpty()) {
                Pair minR = right.pollFirst();
                left.add(minR);
                sumL += minR.val;
            }
        } else {
            right.remove(p);
        }
    }
}

"""
Problem: LeetCode 3013 - Divide an Array Into Subarrays With Minimum Cost II
Difficulty: Hard
Language: Python
Author: CodeByTushu
"""
from typing import List
from sortedcontainers import SortedList


class Solution:
    def minimumCost(self, nums: List[int], k: int, dist: int) -> int:
        n = len(nums)
        m = k - 1

        # Two sorted lists to maintain the m smallest elements dynamically:
        # left: contains the m smallest elements currently in the window
        # right: contains the remaining elements currently in the window
        left = SortedList()
        right = SortedList()
        sum_l = 0

        def add(val: int, idx: int):
            nonlocal sum_l
            p = (val, idx)
            if len(left) < m:
                left.add(p)
                sum_l += val
            else:
                if p < left[-1]:
                    max_l = left.pop()
                    sum_l -= max_l[0]
                    right.add(max_l)

                    left.add(p)
                    sum_l += val
                else:
                    right.add(p)

        def remove(val: int, idx: int):
            nonlocal sum_l
            p = (val, idx)
            if p in left:
                left.remove(p)
                sum_l -= val
                if len(right) > 0:
                    min_r = right.pop(0)
                    left.add(min_r)
                    sum_l += min_r[0]
            else:
                right.remove(p)

        # Initialize the first window of size dist + 1: indices 1 to min(n - 1, 1 + dist)
        for i in range(1, min(n, 2 + dist)):
            add(nums[i], i)

        min_sum = sum_l

        # Slide the window across nums[1 ... n - 1]
        for i in range(2 + dist, n):
            add(nums[i], i)
            remove(nums[i - dist - 1], i - dist - 1)
            min_sum = min(min_sum, sum_l)

        return nums[0] + min_sum

/**
 * Problem: LeetCode 3013 - Divide an Array Into Subarrays With Minimum Cost II
 * Difficulty: Hard
 * Language: C++
 * Author: CodeByTushu
 */
#include <vector>
#include <set>
#include <algorithm>

using namespace std;

class Solution {
    set<pair<long long, int>> left;
    set<pair<long long, int>> right;
    long long sumL = 0;
    int m = 0;

    void add(pair<long long, int> p) {
        if (left.size() < (size_t)m) {
            left.insert(p);
            sumL += p.first;
        } else {
            if (p < *left.rbegin()) {
                auto maxIt = prev(left.end());
                pair<long long, int> maxL = *maxIt;
                left.erase(maxIt);
                sumL -= maxL.first;
                right.insert(maxL);

                left.insert(p);
                sumL += p.first;
            } else {
                right.insert(p);
            }
        }
    }

    void remove(pair<long long, int> p) {
        auto it = left.find(p);
        if (it != left.end()) {
            left.erase(it);
            sumL -= p.first;
            if (!right.empty()) {
                auto minIt = right.begin();
                pair<long long, int> minR = *minIt;
                right.erase(minIt);
                left.insert(minR);
                sumL += minR.first;
            }
        } else {
            right.erase(p);
        }
    }

public:
    long long minimumCost(vector<int>& nums, int k, int dist) {
        int n = nums.size();
        m = k - 1;
        left.clear();
        right.clear();
        sumL = 0;

        for (int i = 1; i <= min(n - 1, 1 + dist); ++i) {
            add({nums[i], i});
        }

        long long minSum = sumL;

        for (int i = 2 + dist; i < n; ++i) {
            add({nums[i], i});
            remove({nums[i - dist - 1], i - dist - 1});
            minSum = min(minSum, sumL);
        }

        return (long long)nums[0] + minSum;
    }
};

/**
 * Problem: LeetCode 3013 - Divide an Array Into Subarrays With Minimum Cost II
 * Difficulty: Hard
 * Language: JavaScript
 * Author: CodeByTushu
 */
/**
 * @param {number[]} nums
 * @param {number} k
 * @param {number} dist
 * @return {number}
 */
var minimumCost = function(nums, k, dist) {
    const n = nums.length;
    const m = k - 1;

    // Coordinate compression over nums[1 ... n - 1]
    const vals = Array.from(new Set(nums.slice(1))).sort((a, b) => a - b);
    const U = vals.length;

    function getRank(x) {
        let l = 0, r = U - 1;
        while (l <= r) {
            const mid = (l + r) >> 1;
            if (vals[mid] === x) return mid + 1;
            if (vals[mid] < x) l = mid + 1;
            else r = mid - 1;
        }
        return l + 1;
    }

    const bitCount = new Int32Array(U + 1);
    const bitSum = new Float64Array(U + 1);

    function update(rank, val, delta) {
        for (let i = rank; i <= U; i += i & -i) {
            bitCount[i] += delta;
            bitSum[i] += delta * val;
        }
    }

    function querySumOfSmallest(count) {
        let idx = 0;
        let currentCount = 0;
        let step = 1;
        while ((step << 1) <= U) step <<= 1;

        for (; step > 0; step >>= 1) {
            if (idx + step <= U && currentCount + bitCount[idx + step] < count) {
                idx += step;
                currentCount += bitCount[idx];
            }
        }

        let sum = 0;
        for (let i = idx; i > 0; i -= i & -i) {
            sum += bitSum[i];
        }

        const rem = count - currentCount;
        if (rem > 0 && idx < U) {
            sum += rem * vals[idx];
        }
        return sum;
    }

    // Initialize first window of size dist + 1: indices 1 to min(n - 1, 1 + dist)
    for (let i = 1; i <= Math.min(n - 1, 1 + dist); i++) {
        update(getRank(nums[i]), nums[i], 1);
    }

    let minSum = querySumOfSmallest(m);

    // Slide window across nums[1 ... n - 1]
    for (let i = 2 + dist; i < n; i++) {
        update(getRank(nums[i]), nums[i], 1);
        update(getRank(nums[i - dist - 1]), nums[i - dist - 1], -1);
        const cur = querySumOfSmallest(m);
        if (cur < minSum) minSum = cur;
    }

    return nums[0] + minSum;
};