3820. Pythagorean Distance Nodes in a Tree

You are given an integer n and an undirected tree with n nodes numbered from 0 to n - 1. The tree is represented by a 2D array edges of length n - 1, where edges[i] = [u_i, v_i] indicates an undirected edge between u_i and v_i.

You are also given three distinct target nodes x, y, and z.

For any node u in the tree:

• Let dx be the distance from u to node x
• Let dy be the distance from u to node y
• Let dz be the distance from u to node z

The node u is called special if the three distances form a Pythagorean Triplet.

Return an integer denoting the number of special nodes in the tree.

A Pythagorean triplet consists of three integers a, b, and c which, when sorted in ascending order, satisfy a^2 + b^2 = c^2.

The distance between two nodes in a tree is the number of edges on the unique path between them.

const int MAX_N = 101010;
vector<pair<int,int>> adj[MAX_N];
long long level[MAX_N], LCA[MAX_N][22], dep[MAX_N];
void dfs(long long u, long long lvl, long long par) {
    level[u] = lvl;
    LCA[u][0] = par;
    for(int i = 1; i < 22; i++) {
        LCA[u][i] = LCA[LCA[u][i-1]][i-1];
    }
    for(auto& [v,w] : adj[u]) {
        if(v == par) continue;
        dep[v] = dep[u] + w;
        dfs(v, lvl + 1, u);
    }
}
long long lcaQuery(long long u, long long v) {
    if(level[u] < level[v]) swap(u, v);
    long long diff = level[u] - level[v];
    for(long long i = 0; diff; i++, diff /= 2) {
        if(diff & 1) u = LCA[u][i];
    }
    if(u != v) {
        for(int i = 21; i >= 0; i--) {
            if(LCA[u][i] == LCA[v][i]) continue;
            u = LCA[u][i];
            v = LCA[v][i];
        }
        u = LCA[u][0];
    }
    return u;
}
long long distance(long long u, long long v) {
    long long lca = lcaQuery(u,v);
    return dep[u] + dep[v] - 2 * dep[lca];
}
class Solution {
public:
    int specialNodes(int n, vector<vector<int>>& edges, int x, int y, int z) {
        memset(LCA,0,sizeof LCA);
        for(int i = 1; i <= n; i++) adj[i].clear();
        for(auto& e : edges) {
            int u = e[0] + 1, v = e[1] + 1, w = 1;
            adj[u].push_back({v,w});
            adj[v].push_back({u,w});
        }
        x += 1, y += 1, z += 1;
        dfs(1,0,0);
        int res = 0;
        for(int i = 1; i <= n; i++) {
            long long dx = distance(i,x), dy = distance(i,y), dz = distance(i,z);
            vector<long long> d{dx,dy,dz};
            sort(begin(d), end(d));
            if(d[0] * d[0] + d[1] * d[1] == d[2] * d[2]) res++;
        }
        return res;
    }
};