Codeforces
CF Step
Youtube Linkedin Discord Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Back to homepage

Code : Find X-Sum of All K-Long Subarrays II

class Solution {
  public:
    vector<long long> findXSum(vector<int> &nums, int k, int x);
};

vector<long long> Solution::findXSum(vector<int> &a, int k, int x) {
    int n = a.size();
    vector<long long> ans;
    set<pair<int, int>> selected, waitlist;
    // selected is a min heap, waitlist is a max heap.

    long long skill_sum = 0;

    auto remove_candidate = [&](pair<int, int> cand) {
        // If this was a waitlisted candidate, they can simply back out of the
        // interview process.
        if (waitlist.find(cand) != waitlist.end()) {
            waitlist.erase(cand);
        }

        // If this candidate was already selected and they want to back out of
        // the interview process, pick the most proficient waitilisted
        // candidate.
        if (selected.find(cand) != selected.end()) {
            selected.erase(cand);
            skill_sum -= 1LL * cand.first * cand.second;
            if (!waitlist.empty()) {
                auto new_selection = *waitlist.rbegin();
                waitlist.erase(new_selection);
                selected.insert(new_selection);
                skill_sum += 1LL * new_selection.first * new_selection.second;
            }
        }
    };

    auto add_candidate = [&](pair<int, int> cand) {
        // A new candidate has arrived. Select them blindly, then reject the
        // least proficient selected candidate.
        selected.insert(cand);
        skill_sum += 1LL * cand.first * cand.second;

        // If we have selected more than k candidates, move one candidate to
        // the waiting list.
        if (selected.size() > x) {
            // Transfer the least eligible selected candidate to the waiting
            // list.
            auto rejection = *selected.begin();
            selected.erase(rejection);
            skill_sum -= 1LL * rejection.first * rejection.second;
            waitlist.insert(rejection);
        }
    };

    map<int, int> freq;
    for (int i = 0; i < n; i++) {
        pair<int, int> cand = {freq[a[i]], a[i]};
        // If the candidate already exists, remove it. We will reapply.
        if (freq[a[i]] > 0) {
            remove_candidate(cand);
        }
        freq[a[i]]++;
        cand = {freq[a[i]], a[i]};
        add_candidate(cand);

        if (i >= k) {
            // Decrease the skillset of the oldest candidate.
            // To do that, we remove the previous application and apply fresh
            // for the job.
            int out = i - k;
            cand = {freq[a[out]], a[out]};
            remove_candidate(cand);
            freq[a[out]]--;

            if (freq[a[out]] > 0) {
                cand = {freq[a[out]], a[out]};
                add_candidate(cand);
            }
        }

        if (i >= k - 1) {
            ans.push_back(skill_sum);
        }
    }
    return ans;
}