1804G - Flow Control - CodeForces Solution

#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <vector>

// #define TIME_START auto time_start = std::chrono::system_clock::now();
// #define TIME_END                                                                                                                         \
//     int64_t elapsed_time = std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now() - time_start).count(); \
//     std::cout << "elapsed time: " << elapsed_time << "\n";

int64_t sum(int64_t duration, int64_t total_speed, int64_t size) {
    int64_t start_speed = total_speed;
    int64_t end_speed = total_speed + (duration - 1) * size;

    return (start_speed + end_speed) * duration / 2;
}

struct Routine {
    int64_t bandwidth = 0;

    std::map<int, int64_t> residual_map;
    std::map<int64_t, std::set<int>> residual_set;
    int64_t current_time = 0;
    int64_t common_speed = 0;

    int64_t total_speed = 0;
    int64_t answer = 0;

    int64_t debug_decrement = 0;
    int64_t debug_process_until = 0;
    int64_t debug_process_timestep = 0;
    int64_t debug_process_multiple_residual = 0;
    int64_t debug_process_single_residual = 0;
    int64_t debug_process_double_residual = 0;

    std::map<std::pair<int64_t, int64_t>, int64_t> debug_double;

    void set_bandwidth(int64_t b) {
        bandwidth = b;
    }

    void insert(int idx, int64_t speed) {
        residual_map[idx] = speed - common_speed;
        residual_set[speed - common_speed].insert(idx);
        total_speed += speed;
    }

    void erase(int idx) {
        int64_t residual = residual_map[idx];
        residual_set[residual].erase(idx);
        if (residual_set[residual].empty()) {
            residual_set.erase(residual);
        }
        residual_map.erase(idx);
        total_speed -= residual + common_speed;
    }

    void process_timestep() {
        ++debug_process_timestep;
        if (total_speed <= bandwidth) {
            answer += total_speed;
            total_speed += residual_map.size();
            common_speed += 1;
        } else {
            std::map<int64_t, std::set<int>> new_residual_set;
            int64_t new_common_speed = common_speed / 2;

            for (auto& p : residual_set) {
                int64_t residual = p.first;
                int64_t speed = residual + common_speed;

                int64_t new_speed = speed / 2;
                int64_t new_residual = new_speed - new_common_speed;

                total_speed += (new_speed - speed) * p.second.size();

                if (new_residual != residual) {
                    for (int idx : p.second) {
                        residual_map[idx] = new_residual;
                        ++debug_decrement;
                    }
                }

                if (new_residual_set.count(new_residual)) {
                    if (p.second.size() < new_residual_set[new_residual].size()) {
                        for (int idx : p.second) {
                            new_residual_set[new_residual].insert(idx);
                        }
                    } else {
                        for (int idx : new_residual_set[new_residual]) {
                            p.second.insert(idx);
                        }
                        std::swap(p.second, new_residual_set[new_residual]);
                    }
                    p.second.clear();
                } else {
                    std::swap(new_residual_set[new_residual], p.second);
                }
            }

            common_speed = new_common_speed;
            std::swap(residual_set, new_residual_set);
        }

        ++current_time;
    }

    void process_multiple_residual(int64_t end_time) {
        ++debug_process_multiple_residual;

        if (residual_set.size() > 3) {
            process_until_overflow(end_time);
            return;
        }

        std::set<std::pair<int64_t, int64_t>> set;
        int64_t cs = common_speed;
        for (auto& p : residual_set) {
            set.insert({p.first, p.second.size()});
        }

        process_until_overflow(end_time);

        if (current_time == end_time) {
            return;
        }

        std::set<std::pair<int64_t, int64_t>> new_set;
        int64_t new_cs = common_speed;
        for (auto& p : residual_set) {
            new_set.insert({p.first, p.second.size()});
        }

        if (new_set != set || new_cs != cs) {
            return;
        }

        // total_speed + (duration - 1) * residual_map.size() <= bandwidth
        int64_t duration = (bandwidth - total_speed) / residual_map.size() + 1;
        int64_t add_answer = sum(duration, total_speed, residual_map.size());

        int64_t cnt = (end_time - current_time) / (duration + 1);
        current_time += (duration + 1) * cnt;
        answer += add_answer * cnt;
    }

    void process_until_overflow(int64_t end_time) {
        if (total_speed > bandwidth) {
            process_timestep();
            return;
        }
        // total_speed + (duration - 1) * residual_map.size() <= bandwidth
        int64_t duration = (bandwidth - total_speed) / residual_map.size() + 1;
        duration = std::min(duration, end_time - current_time);

        answer += sum(duration, total_speed, residual_map.size());
        total_speed += duration * residual_map.size();
        common_speed += duration;
        current_time += duration;

        if (current_time < end_time) {
            process_timestep();
        }
    }

    void process_until(int64_t end_time) {
        ++debug_process_until;
        while (current_time < end_time) {
            if (residual_set.size() == 0) {
                common_speed = 0;
                current_time = end_time;
            } else {
                process_multiple_residual(end_time);
            }
        }
    }
};

void solve() {
    int n = 0;
    int64_t bandwidth = 0;
    std::cin >> n >> bandwidth;

    std::map<int64_t, std::vector<std::pair<int, bool>>> timeline;
    std::vector<int64_t> speed(n);
    for (int i = 0; i < n; ++i) {
        int64_t s = 0, f = 0, d = 0;
        std::cin >> s >> f >> d;
        timeline[s].push_back({i, true});
        timeline[f + 1].push_back({i, false});
        speed[i] = d;
    }

    Routine routine;
    routine.set_bandwidth(bandwidth);

    for (auto& p : timeline) {
        routine.process_until(p.first);

        for (auto event : p.second) {
            if (event.second) {
                routine.insert(event.first, speed[event.first]);
            } else {
                routine.erase(event.first);
            }
        }

        routine.process_timestep();
    }

    std::cout << routine.answer << "\n";
    // std::cout << routine.debug_decrement / n << "\n";
    // std::cout << routine.debug_process_until << "\n";
    // std::cout << routine.debug_process_timestep << "\n";
    // std::cout << routine.debug_process_multiple_residual << "\n";
    // std::cout << routine.debug_process_single_residual << "\n";
    // std::cout << routine.debug_process_double_residual << "\n";
}

int main() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);

#ifndef ONLINE_JUDGE
    std::freopen("input.txt", "r", stdin);
#endif

    solve();

    return 0;
}