|
| 1 | +/** |
| 2 | + * @param {number[][]} edges |
| 3 | + * @param {number[][]} queries |
| 4 | + * @return {number[]} |
| 5 | + */ |
| 6 | +var assignEdgeWeights = function (edges, queries) { |
| 7 | + const n = edges.length + 1, |
| 8 | + graph = Array(n) |
| 9 | + .fill(0) |
| 10 | + .map(() => []) |
| 11 | + for (let [u, v] of edges) { |
| 12 | + graph[u - 1].push(v - 1) |
| 13 | + graph[v - 1].push(u - 1) |
| 14 | + } |
| 15 | + const directParent = Array(n), |
| 16 | + depth = Array(n) |
| 17 | + dfs(0, -1, 0) |
| 18 | + const lca = new LCA(n, directParent, depth) |
| 19 | + const answer = [], |
| 20 | + MOD = 1000000007n |
| 21 | + for (let [u, v] of queries) { |
| 22 | + const pathLen = lca.getDist(u - 1, v - 1) |
| 23 | + const ways = pathLen === 0 ? 0n : 2n ** BigInt(pathLen - 1) |
| 24 | + answer.push(Number(ways % MOD)) |
| 25 | + } |
| 26 | + return answer |
| 27 | + |
| 28 | + function dfs(node, parent, currDepth) { |
| 29 | + directParent[node] = parent |
| 30 | + depth[node] = currDepth |
| 31 | + for (let nei of graph[node]) { |
| 32 | + if (nei === parent) continue |
| 33 | + dfs(nei, node, currDepth + 1) |
| 34 | + } |
| 35 | + } |
| 36 | +} |
| 37 | + |
| 38 | +class LCA { |
| 39 | + constructor(n, parent, depths) { |
| 40 | + this.maxDepth = Math.ceil(Math.log2(n)) |
| 41 | + this.p = Array(this.maxDepth + 1) |
| 42 | + .fill(0) |
| 43 | + .map(() => Array(n).fill(-1)) // parents |
| 44 | + this.depths = depths |
| 45 | + for (let node = 0; node < n; node++) { |
| 46 | + this.p[0][node] = parent[node] |
| 47 | + } |
| 48 | + for (let pow2 = 1; pow2 <= this.maxDepth; pow2++) { |
| 49 | + for (let node = 0; node < n; node++) { |
| 50 | + const halfParent = this.p[pow2 - 1][node] |
| 51 | + if (halfParent !== -1) { |
| 52 | + this.p[pow2][node] = this.p[pow2 - 1][halfParent] |
| 53 | + } |
| 54 | + } |
| 55 | + } |
| 56 | + } |
| 57 | + getLCA(a, b) { |
| 58 | + if (this.depths[a] > this.depths[b]) { |
| 59 | + let temp = a |
| 60 | + a = b |
| 61 | + b = temp |
| 62 | + } |
| 63 | + let depthDiff = this.depths[b] - this.depths[a] |
| 64 | + for (let i = 0; i <= this.maxDepth; i++) { |
| 65 | + if ((depthDiff >> i) & 1) { |
| 66 | + b = this.p[i][b] // move b up to the 2^ith parent |
| 67 | + } |
| 68 | + } |
| 69 | + if (a === b) return a |
| 70 | + |
| 71 | + // move both nodes up by 2^ith levels if the 2^ith parents are not equal |
| 72 | + for (let i = this.maxDepth; i >= 0; i--) { |
| 73 | + // this decrements so that we can jump the nodes up incrementally |
| 74 | + if (this.p[i][a] !== this.p[i][b]) { |
| 75 | + // if 2^ith parents of both nodes are not equal, we can safely both move up |
| 76 | + a = this.p[i][a] |
| 77 | + b = this.p[i][b] |
| 78 | + } |
| 79 | + } |
| 80 | + return this.p[0][a] |
| 81 | + } |
| 82 | + getDist(a, b) { |
| 83 | + const lca = this.getLCA(a, b) |
| 84 | + const depthA = this.depths[a] - this.depths[lca] |
| 85 | + const depthB = this.depths[b] - this.depths[lca] |
| 86 | + return depthA + depthB |
| 87 | + } |
| 88 | +} |
0 commit comments