[Codeforces] E. Subset Sums

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : E. Subset Sums
Source            : Codeforces
Category          : Data Structure
Algorithm         : Heavy Light Technique
Verdict           : Accepted

#include "bits/stdc++.h"  
  
using namespace std;  
  
#define ll           long long int   
#define endl         '\n'  
  
static const int maxn = 1e5 + 5;  
static const int maxb = 320;  
  
vector <int> light_set[maxn];  
vector <int> heavy_set[maxb];  
int light_heavy_intersection[maxn][maxb];  
int heavy_heavy_intersection[maxb][maxb];  
int num_ele[maxn];  
int mapper[maxn];  
long long lazy[maxn];  
long long ans_heavy[maxn];  
long long arr[maxn];  
  
signed main()  
{  
    ios_base::sync_with_stdio(false);  
    cin.tie(nullptr);  
      
    int n, m, q;  
    cin >> n >> m >> q;  
    for (int i = 1; i <= n; i++) cin >> arr[i];  
    int light = 0;  
    int heavy = 0;  
    for (int i = 1; i <= m; i++)  
    {  
        cin >> num_ele[i];  
        if (num_ele[i] > maxb)  
        {  
            mapper[i] = ++heavy;  
            long long sum = 0;  
            heavy_set[heavy].resize(num_ele[i]);  
            for (int &x : heavy_set[heavy])   
            {  
                cin >> x;  
                sum += arr[x];        
            }  
            sort(heavy_set[heavy].begin(), heavy_set[heavy].end());  
            ans_heavy[heavy] = sum;  
        }  
        else   
        {  
            mapper[i] = ++light;  
            light_set[light].resize(num_ele[i]);  
            for (int &x : light_set[light]) cin >> x;  
            sort(light_set[light].begin(), light_set[light].end());  
        }  
    }  
    auto get = [&](vector <int> &vec, int key)  
    {  
        auto fnd = lower_bound(vec.begin(), vec.end(), key);  
        return fnd != vec.end() && *fnd == key;   
    };  
    for (int lgt = 1; lgt <= light; lgt++)  
    {  
        for (int x : light_set[lgt])  
        {  
            for (int hvy = 1; hvy <= heavy; hvy++)   
                light_heavy_intersection[lgt][hvy] += get(heavy_set[hvy], x);  
        }  
    }  
    for (int hvy = 1; hvy <= heavy; hvy++)  
    {  
        for (int x : heavy_set[hvy])  
        {  
            for (int i = 1; i <= heavy; i++) if (i != hvy)  
                heavy_heavy_intersection[hvy][i] += get(heavy_set[i], x);    
        }  
    }  
    while (q--)  
    {  
        char type;  
        cin >> type;  
        if (type == '?')  
        {  
            int x;  
            cin >> x;  
            if (num_ele[x] > maxb)  
            {  
                int hvy = mapper[x];  
                long long res = ans_heavy[hvy] + 1LL * num_ele[x] * lazy[hvy];  
                for (int i = 1; i <= heavy; i++) if (hvy != i)  
                    res += heavy_heavy_intersection[hvy][i] * lazy[i];  
                cout << res << endl;  
            }  
            else  
            {  
                int lgt = mapper[x];  
                long long res = 0;  
                for (int p : light_set[lgt]) res += arr[p];  
                for (int z = 1; z <= heavy; z++) res += (light_heavy_intersection[lgt][z] * lazy[z]);  
                cout << res << endl;  
            }  
        }  
        else   
        {  
            int x;  
            long long val;  
            cin >> x >> val;  
            if (num_ele[x] > maxb)   
            {  
                int hvy = mapper[x];  
                lazy[hvy] += val;  
            }  
            else   
            {  
                int lgt = mapper[x];  
                for (int p : light_set[lgt]) arr[p] += val;  
                for (int hvy = 1; hvy <= heavy; hvy++) ans_heavy[hvy] += (light_heavy_intersection[lgt][hvy] * val);  
            }  
        }  
    }  
}