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


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