LCA.code-snippets

{
	"LCA": {
		"prefix": "LCA",
		"body": [
		  "template < typename T = int > struct LCA {",
		  "    ",
		  "    struct Edge {",
		  "",
		  "        ll v, w;",
		  "",
		  "        Edge(ll V = 0, ll W = 0) : v(V), w(W) {}",
		  "",
		  "        bool operator < (const Edge &rhs) const {",
		  "            return w < rhs.w;",
		  "        }",
		  "",
		  "    };",
		  "",
		  "    int N, LOG;",
		  "    vector < vector < T > > anc, cost;",
		  "    vector < vector < Edge > > adj;",
		  "    vector < int > dep;",
		  "    ",
		  "    LCA(int n){",
		  "        N = n + 10, LOG = 0;",
		  "        while((1 << LOG) <= n) LOG++;",
		  "        dep = vector < int > (N);",
		  "        adj = vector < vector < Edge > > (N);",
		  "        anc = cost = vector < vector < T > > (N, vector < T > (LOG));",
		  "    }",
		  "    ",
		  "    void add_edge(int u, int v, T w){",
		  "        adj[u].push_back(Edge(v, w));",
		  "        adj[v].push_back(Edge(u, w));",
		  "    }",
		  "",
		  "    void build_adj(int edges){",
		  "        for(int i = 0, u, v, w; i < edges && cin >> u >> v >> w; i++)",
		  "            add_edge(u, v, w);",
		  "    }",
		  "",
		  "    T operation(T a, T b){",
		  "        return a + b;",
		  "    }",
		  "    ",
		  "    void dfs(int u, int p = 0){",
		  "        for(auto& [v, w] : adj[u]){",
		  "            if(v == p) continue;",
		  "            dep[v] = dep[u] + 1, anc[v][0] = u, cost[v][0] = w;",
		  "            for(int bit = 1; bit < LOG; bit++){",
		  "                anc[v][bit] = anc[anc[v][bit - 1]][bit - 1];",
		  "                cost[v][bit] = operation(cost[v][bit - 1], cost[anc[v][bit - 1]][bit - 1]);",
		  "            }",
		  "            dfs(v, u);",
		  "        }",
		  "    }",
		  "    ",
		  "    int kth_ancestor(int u, int k){",
		  "        if(dep[u] < k) ",
		  "            return -1;",
		  "        for(int bit = LOG - 1; bit >= 0; bit--)",
		  "            if(k & (1 << bit))",
		  "                u = anc[u][bit];",
		  "        return u;",
		  "    }",
		  "    ",
		  "    int get_lca(int u, int v){",
		  "        if(dep[u] < dep[v])",
		  "            swap(u, v);",
		  "        u = kth_ancestor(u, dep[u] - dep[v]);",
		  "        if(u == v)",
		  "            return u;",
		  "        for(int bit = LOG - 1; bit >= 0; bit--)",
		  "            if(anc[u][bit] != anc[v][bit])",
		  "                u = anc[u][bit], v = anc[v][bit];",
		  "        return anc[u][0];",
		  "    }",
		  "    ",
		  "    T get_cost(int u, int dist){",
		  "        if(dep[u] < dist) return -1;",
		  "        T ans = 0;",
		  "        for(int bit = 0; bit < LOG; bit++)",
		  "            if(dist & (1 << bit))",
		  "                ans = operation(ans, cost[u][bit]), u = anc[u][bit];",
		  "        return ans;",
		  "    }",
		  "    ",
		  "    T query(int u, int v){",
		  "        int lca = get_lca(u, v);",
		  "        return operation(get_cost(u, dep[u] - dep[lca]), get_cost(v, dep[v] - dep[lca]));",
		  "    }",
		  "",
		  "};"
		],
		"description": "LCA"
	}
}