Create 3321. Find X-Sum of All K-Long Subarrays II (#927)

Author: Chayandas07

3321. Find X-Sum of All K-Long Subarrays II

@@ -0,0 +1,78 @@
+class Solution {
+public:
+    using P = pair<int,int>;
+    struct Cmp {
+        bool operator()(const P& a, const P& b) const {
+            if (a.first != b.first) return a.first > b.first;
+            return a.second > b.second;
+        }
+    };
+
+    vector<long long> findXSum(vector<int>& nums, int k, int x) {
+        int n = (int)nums.size();
+        vector<long long> ans(n - k + 1);
+
+        unordered_map<int,int> cnt;
+        cnt.reserve(n * 1);
+
+        set<P, Cmp> top, rest;
+        long long topSum = 0;
+
+        auto pull = [&](int v, int f){
+            P key{f, v};
+            auto it = top.find(key);
+            if (it != top.end()){
+                topSum -= 1LL * v * f;
+                top.erase(it);
+            } else {
+                auto jt = rest.find(key);
+                if (jt != rest.end()) rest.erase(jt);
+            }
+        };
+        auto pushToTop = [&](int v, int f){
+            top.insert({f, v});
+            topSum += 1LL * v * f;
+        };
+
+        auto insertVal = [&](int v){
+            int f = cnt[v];
+            if (f) pull(v, f);
+            ++f; cnt[v] = f;
+            pushToTop(v, f);
+            if ((int)top.size() > x){
+                auto it = prev(top.end());
+                topSum -= 1LL * it->first * it->second;
+                rest.insert(*it);
+                top.erase(it);
+            }
+        };
+
+        auto eraseVal = [&](int v){
+            auto itc = cnt.find(v);
+            if (itc == cnt.end()) return;
+            int f = itc->second;
+            pull(v, f);
+            if (--f == 0){
+                cnt.erase(itc);
+            } else {
+                cnt[v] = f;
+                rest.insert({f, v});
+            }
+            if ((int)top.size() < x && !rest.empty()){
+                auto best = rest.begin();
+                topSum += 1LL * best->first * best->second;
+                top.insert(*best);
+                rest.erase(best);
+            }
+        };
+
+        for (int i = 0; i < k; ++i) insertVal(nums[i]);
+        ans[0] = topSum;
+        for (int i = k; i < n; ++i){
+            eraseVal(nums[i-k]);
+            insertVal(nums[i]);
+            ans[i - k + 1] = topSum;
+        }
+        return ans;
+    }
+};