HackerEarth Infinite K-tree problem solution

In this HackerEarth Infinite K-tree problem solution, You're given a K-ary infinite tree rooted at a vertex numbered 1. All its edges are weighted 1 initially.

Any node X will have exactly K children numbered as:

[K * X + 0, K * X + 1, K * X + 2, K * X + 3, K * X + 4,........K * X + (K - 1)]

You are given Q queries to answer which will be of the following two types:
1. 1uv: Print the shortest distance between nodes u and v.
2. 2uvw: Increase the weight of all edges on the shortest path between u and v by w.

HackerEarth Infinite K-tree problem solution.

`#include <iostream>#include <map>#include <assert.h>using namespace std;#define int long longmap < pair < int, int >, int > adj;int find_depth( int u, int k ) {    int depth = 0;    while ( u > 0 ) {        u = u / k;        depth = depth + 1;    }    return depth - 1;}int dist( int u, int v, int k ) {    int dist = 0;    int depth_u = find_depth( u, k );    int depth_v = find_depth( v, k );    if ( depth_u < depth_v ) {    swap ( u, v );    swap ( depth_u, depth_v );    }    while( depth_u != depth_v ) {        if ( adj.count( { u, u / k } ) ) {        dist = dist + adj[ { u, u / k } ];    } else {        dist = dist + 1;    }    depth_u = depth_u - 1;    u = u / k;    }    while ( u != v ) {    if ( adj.count( { u, u / k } ) ) {        dist = dist + adj [ { u, u / k } ];    } else {        dist = dist + 1;    }    if ( adj.count( { v, v / k } ) ) {                dist = dist + adj [ { v, v / k } ];        } else {                dist = dist + 1;        }        u = u / k;        v = v / k;    }    return dist;}void add_weight( int vertex, int parent, int w ) {    if ( !adj.count ( { vertex, parent } ) ) {        adj[ { vertex, parent } ] = 1;    }    adj[ { vertex, parent } ] = adj[ { vertex, parent } ] + w;}void increase_weights ( int u, int v, int w, int k ) {    int depth_u = find_depth( u, k );    int depth_v = find_depth( v, k );    if ( depth_u < depth_v ) {     swap ( u, v );    swap ( depth_u, depth_v );    }    while( depth_u != depth_v ) {        add_weight( u, u / k, w );    depth_u = depth_u - 1;    u = u / k;    }    while ( u != v ) {    add_weight( u, u / k, w );    add_weight( v, v / k, w );        u = u / k;        v = v / k;    }}signed main() {    ios::sync_with_stdio(0);    cin.tie(0);    cout.tie(0);    int k, q, x, u, v, w;    cin >> k >> q;    while ( q-- ) {        cin >> x;        if ( x == 1 ) {            cin >> u >> v;            cout << dist( u, v, k ) << "\n";        } else {            cin >> u >> v >> w;            increase_weights( u, v, w, k );        }    }}`

Second solution

`from collections import defaultdictk, q = map(int, input().split())def path(v):    ans = [v]    while v != 1:        v //= k        ans += [v]    return ansdef lca(v, u):    s = set(path(v))    for x in path(u):        if x in s:            return xweight_to_par = defaultdict(lambda: 1)while q > 0:    q -= 1    query = map(int, input().split())    if next(query) == 1:        v, u = query        ans = 0        p = lca(v, u)        while v != p:            ans += weight_to_par[v]            v //= k        while u != p:            ans += weight_to_par[u]            u //= k        print(ans)    else:        v, u, w = query        p = lca(v, u)        while v != p:            weight_to_par[v] += w            v //= k        while u != p:            weight_to_par[u] += w            u //= k`