HackerEarth Sum of shortest paths problem solution

In this HackerEarth Sum of shortest paths problem solution Consider that D(G,u,v) is defined as the number of edges on the shortest path from u to v in a graph G.

You are given a tree T with N vertices and Q queries of the following type:

• If we add edge (a,b) to the tree T, obtaining graph G1, then what is the value of Sigma(1<=u<=v<=N) D(G1,u,v)?

Note: The queries are independent in nature.

HackerEarth Sum of shortest paths problem solution.

#include "bits/stdc++.h"
using namespace std;
#define fi first
#define se second
#define ll long long
#define dbg(v) cerr<<#v<<" = "<<v<<'\n'
#define vi vector<int>
#define vl vector <ll>
#define pii pair<int,int>
#define vii vector < pii >
#define mp make_pair
#define db long double
#define pb push_back
#define all(s) s.begin(),s.end()
template < class P , class Q > ostream& operator<<(ostream& stream, pair < P , Q > v){ stream << "(" << v.fi << ',' << v.se << ")"; return stream;}
template < class T > ostream& operator<<(ostream& stream, const vector<T> v){ stream << "[ "; for (int i=0; i<(int)v.size(); i++) stream << v[i] << " "; stream << "]"; return stream;}
template < class T > T smin(T &a,T b) {if (a > b) a = b;return a;}
template < class T > T smax(T &a,T b) {if (a < b) a = b;return a;}
const int N = 1e6 + 5;
vi g[N];
int h[N];
ll dp[N];
int F[20][N];
int n;
ll answer = 0;
void dfs(int node,int prev) {
    h[node] = h[prev] + 1;
    F[0][node] = prev;
    dp[node] = 1;
    for (auto it : g[node])
        if (it != prev) {
            dfs(it,node);
            dp[node] += dp[it];
        }
    if (prev)
        answer += dp[node] * (n - dp[node]);
}
int shift(int x,int y) {
    for (int k = 0;k < 20;++k)
        if ((y >> k) & 1)
            x = F[k][x];
    return x;
}
int lca(int x,int y) {
    if (h[x] > h[y])
        x = shift(x,h[x] - h[y]);
    else
        y = shift(y,h[y] - h[x]);
    for (int k = 19;k >= 0;--k)
        if (F[k][x] != F[k][y])
            x = F[k][x],y = F[k][y];
    if (x == y)
        return x;
    return F[0][x];
}
ll s[N];
ll sum1[N];
ll sum2[N];
ll sum3[N];
ll go(vl v) {
    int sz = v.size();
    for (int i = 1;i <= sz;++i)
        s[i] = v[i - 1];
    for (int i = 1;i <= sz;++i) {
        sum1[i] = sum1[i - 1] + s[i];
        sum2[i] = sum2[i - 1] + 1ll * i * s[i];
        sum3[i] = sum3[i - 1] + 1ll * (sz - i + 1) * s[i];
    }
    ll ans = 0;
    const int d = (sz - 1) / 2;
    for (int i = 1;i <= sz;++i) {
        int nxt = min(sz,i + d);
        ans += s[i] * (sum2[nxt] - sum2[i] - i * (sum1[nxt] - sum1[i]));
        if (i + d < sz)
            ans += s[i] * (sum3[sz] - sum3[nxt] + (i - 1) * (sum1[sz] - sum1[nxt]));
    }

    return ans;
}
int main(void) {
    int q;
    scanf("%d%d",&n,&q);
    for (int i = 1;i < n;++i) {
        int u,v;
        scanf("%d%d",&u,&v);
        g[u].pb(v);
        g[v].pb(u);
    }
    const int Root = 1;
    dfs(Root,0);
    for (int k = 1;k < 20;++k)
        for (int i = 1;i <= n;++i)
            F[k][i] = F[k - 1][F[k - 1][i]];
    ll SUM = 0;
    while (q --) {
        int u,v;
        scanf("%d%d",&u,&v);
        int w = lca(u,v);
        if (h[u] > h[v])
            swap(u,v);
        vl s;
        ll ans = answer;
        if (u == w) {
            int a = v;
            s.pb(dp[a]);
            while (a != w) {
                ans -= dp[a] * (n - dp[a]);
                s.pb(dp[F[0][a]] - dp[a]);
                a = F[0][a];
            }
            s.back() = n - dp[shift(v,h[v] - h[w] - 1)];
        } else {
            int a = v,b = u;
            vl X,Y;
            X.pb(dp[a]);
            while (a != w) {
                ans -= dp[a] * (n - dp[a]);
                X.pb(dp[F[0][a]] - dp[a]);
                a = F[0][a];
            }
            Y.pb(dp[b]);
            while (b != w) {
                ans -= dp[b] * (n - dp[b]);
                Y.pb(dp[F[0][b]] - dp[b]);
                b = F[0][b];
            }
            X.pop_back();
            Y.pop_back();
            reverse(all(Y));
            for (auto it : X)
                s.pb(it);
            s.pb(n - dp[shift(v,h[v] - h[w] - 1)] - dp[shift(u,h[u] - h[w] - 1)]);
            for (auto it : Y)
                s.pb(it);
        }
        ans += go(s);
        cout << ans << '\n';
        SUM += s.size();
    }
    return 0;
}

Second solution

#include"bits/stdc++.h"
using namespace std;
 
#define MAX_N (1<<18)
#define MAX MAX_N
 
int n;
int q;
vector<vector<int> > v;
int dist[MAX_N];
void brute() {
    long long int ans = 0;
    for (int i = 0; i < n; i++) {
        memset(dist, -1, sizeof(dist));
        dist[i] = 0;
        queue<int> q;
        q.push(i);
        while (!q.empty()) {
            int b = q.front();
            q.pop();
            for (int go : v[b]) {
                if (dist[go] == -1) {
                    dist[go] = dist[b] + 1;
                    q.push(go);
                }
            }
        }
        for (int i = 0; i < n; i++) {
            ans += dist[i];
        }
    }
    printf("%lld\n", ans / 2LL);
 
}
#define MAX_LOG 19
int lcc[MAX_LOG][MAX];
int dep[MAX];
inline void dfs(int b, int pr = -1, int d = 0) {
    lcc[0][b] = pr;
    dep[b] = d;
    for (int go : v[b]) {
        if (go == pr)continue;
        dfs(go, b, d + 1);
    }
}
void init() {
    for (int i = 0; i + 1 < MAX_LOG; i++) {
        for (int j = 0; j < n; j++) {
            if (lcc[i][j] == -1) {
                lcc[i + 1][j] = -1;
            }
            else {
                lcc[i + 1][j] = lcc[i][lcc[i][j]];
            }
        }
    }
}
int lca(int a, int b) {
    if (dep[a] < dep[b])swap(a, b);
    for (int i = 0; i < MAX_LOG; i++) {
        if (((dep[a] - dep[b]) >> i) & 1) {
            a = lcc[i][a];
        }
    }
    if (a == b)return a;
    for (int i = MAX_LOG - 1; i >= 0; i--) {
        if (lcc[i][a] != lcc[i][b]) {
            a = lcc[i][a];
            b = lcc[i][b];
        }
    }
    return lcc[0][a];
}
int up[MAX];
int up_sz;
void cyc(int a, int b) {
    up_sz = 0;
    int lc = lca(a, b);
    int sz = 0;
    while (a != lc) {
        up[up_sz++] = a;
        a = lcc[0][a];
    }
    up[up_sz++] = lc;
    int sz2 = up_sz;
    while (b != lc) {
        up[up_sz++] = b;
        b = lcc[0][b];
    }
    reverse(up + sz2, up + up_sz);
    return;
}
struct st {
    long long int ans;
    long long int sz;
    long long int inner;
    st(long long int a_ = 0, long long int b_ = 0, long long int inner_ = 0) {
        ans = a_;
        sz = b_;
        inner = inner_;
    }
    inline st gain() {
        ans += sz;
        inner += ans;
        sz += 1;
        return (*this);
    }
};
inline st merge(st &a, st &b) {
    return  st(a.ans + b.ans, a.sz + b.sz, a.inner + b.inner + (a.ans + a.sz)*b.sz + (b.ans + b.sz)*a.sz);
}
inline st subtract(st &a, st &b) {
    st r;
    r.ans = a.ans - b.ans;
    r.sz = a.sz - b.sz;
    r.inner = a.inner - b.inner;
    r.inner -= (r.ans + r.sz)*b.sz + (b.ans + b.sz)*r.sz;
    return r;
}
struct TreeDP {
    vector<pair<int, int> > g[MAX];
    vector<pair<st, st> >  im[MAX];
    vector<st>  ind[MAX];
    void init(vector<vector<int> > &vv) {
        //g.assign(n, vector<pair<int, int> >());
        //im.assign(n, vector<pair<st, st> >());
        //ind.assign(n, vector<st>());
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < v[i].size(); j++) {
                g[i].push_back(make_pair(v[i][j], 1));
            }
            int pp = g[i].size();
            sort(g[i].begin(), g[i].end());
            auto u = make_pair(st(), st());
            auto z = st();
            im[i].assign(g[i].size(), u);
            ind[i].assign(g[i].size(), z);
        }
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < v[i].size(); j++) {
                int go = g[i][j].first;
                g[i][j].second = lower_bound(g[go].begin(), g[go].end(), make_pair(i, -1)) - g[go].begin();
            }
        }
    }
    inline st dfs(int b, int pr = -1, int opt2 = -1, bool rec = false) {
        if (rec == false) {
            int pp = pr;
            int p2 = opt2;
            if (pr == -1)pr = -1;
            else pr = lower_bound(g[b].begin(), g[b].end(), make_pair(pr, -1)) - g[b].begin();
            if (opt2 == -1)opt2 = -1;
            else opt2 = lower_bound(g[b].begin(), g[b].end(), make_pair(opt2, -1)) - g[b].begin();
        }
        int lef = pr - 1;
        int rig = pr + 1;
        int id = lef;
        while (id >= 0 && im[b][id].first.sz == 0) {
            if(ind[b][id].sz==0){
                ind[b][id] = im[b][id].first = dfs(g[b][id].first, g[b][id].second, -1, true);
            }
            else{
                im[b][id].first=ind[b][id];
            }
            id--;
        }
        for (int j = id + 1; j <= lef; j++) {
            if (j >= 1)im[b][j].first = merge(im[b][j].first, im[b][j - 1].first);
        }
        id = rig;
        while (id < g[b].size() && im[b][id].second.sz == 0) {
            if(ind[b][id].sz==0){
                ind[b][id] = im[b][id].second = dfs(g[b][id].first, g[b][id].second, -1, true);
            }
            else{
                im[b][id].second=ind[b][id];
            }
            id++;
        }
        for (int j = id - 1; j >= rig; j--) {
            if (j + 1 < g[b].size()) {
                im[b][j].second = merge(im[b][j].second, im[b][j + 1].second);
            }
        }
        st ret;
        if (lef >= 0) {
            ret = im[b][lef].first;
        }
        if (rig < g[b].size()) {
            ret = merge(ret, im[b][rig].second);
        }
        if (opt2 != -1) {
            if (ind[b][opt2].sz == 0) {
                ind[b][opt2] = dfs(g[b][opt2].first, g[b][opt2].second, -1, true);
            }
            ret = subtract(ret, ind[b][opt2]);
        }
        ret.gain();
        return ret;
    }
} TD;
set<pair<int, int> > s;
pair<long long int,long long int> temp[MAX];
int main() {
    cin >> n >> q;
    v.assign(n, vector<int>());
    for (int i = 1; i < n; i++) {
        int a, b;
        scanf("%d%d", &a, &b);
        a--;
        b--;
        v[a].push_back(b);
        v[b].push_back(a);
        s.insert(make_pair(a, b));
        s.insert(make_pair(b, a));
    }
    dfs(0);
    TD.init(v);
    init();
    while (q--) {
        int a, b;
        scanf("%d%d", &a, &b);
        a--;
        b--;
        assert(a >= 0 && a < n&&b >= 0 && b < n);
        if (s.count(make_pair(a, b))) {
            long long int ans = TD.dfs(0, -1).inner;
            printf("%lld\n", ans);
            continue;
        }
        cyc(a, b);
        long long int ans = 0;
        int tt=0;
        for (int i = 0; i < up_sz; i++) {
            int nodeA, nodeB;
            if (i == up_sz - 1) {
                nodeA = -1;
            }
            else {
                nodeA = up[i + 1];
            }
            if (i == 0) {
                nodeB = -1;
            }
            else {
                nodeB = up[i - 1];
            }
            st ALL = TD.dfs(up[i], nodeA, nodeB);
            ans += ALL.inner;
            int sz = ALL.sz;
            long long int val = ALL.ans;
            temp[i]=make_pair(val,sz);
            //temp.push_back(make_pair(val, sz));
        }
        queue<pair<long long int, long long int> > dd;
        long long int sz2 = 0;
        long long int cur2 = 0;
        long long int cur1 = 0;
        long long int sz1 = 0;
        int lim = up_sz / 2;
        for (int i = 0; i < up_sz; i++) {
            dd.push(temp[i]);
 
            cur1 += sz1;
            cur2 -= sz2;
            ans += (cur1 + cur2)*temp[i].second + temp[i].first*(sz1 + sz2);
            //cout << (cur1 + cur2)*temp[i].second + temp[i].first*(sz1 + sz2) << endl;
            sz1 += temp[i].second;
            cur1 += temp[i].first;
            if ((int)(dd.size()) - 1 == lim) {
                sz1 -= dd.front().second;
                cur1 -= dd.front().first;
                cur1 -= dd.front().second*lim;
                sz2 += dd.front().second;
                cur2 += dd.front().first;
                cur2 += dd.front().second*(up_sz - lim);
                dd.pop();
            }
        }
        printf("%lld\n", ans);
    }
    return 0;
}