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 long

map < 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 defaultdict

k, q = map(int, input().split())


def path(v):
    ans = [v]
    while v != 1:
        v //= k
        ans += [v]
    return ans


def lca(v, u):
    s = set(path(v))
    for x in path(u):
        if x in s:
            return x


weight_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