diff --git a/Binary_Search/3585.Find-Weighted-Median-Node-in-Tree/3585.Find-Weighted-Median-Node-in-Tree.cpp b/Binary_Search/3585.Find-Weighted-Median-Node-in-Tree/3585.Find-Weighted-Median-Node-in-Tree.cpp new file mode 100644 index 000000000..15e4011a6 --- /dev/null +++ b/Binary_Search/3585.Find-Weighted-Median-Node-in-Tree/3585.Find-Weighted-Median-Node-in-Tree.cpp @@ -0,0 +1,110 @@ +using ll = long long; +const int MAXN = 100000; +const int LOGN = 17; + +class Solution { +public: + vector> adj[MAXN]; + int up[MAXN][LOGN+1]; + int depth[MAXN]; + ll distRoot[MAXN]; + + void dfs(int cur, int parent) + { + up[cur][0] = parent; + for(auto &[v,w]: adj[cur]) + { + if(v == parent) continue; + depth[v] = depth[cur] + 1; + distRoot[v] = distRoot[cur] + w; + dfs(v, cur); + } + } + + int lca(int a, int b) + { + if(depth[a] < depth[b]) swap(a,b); + int diff = depth[a] - depth[b]; + for(int k = 0; k <= LOGN; k++){ + if(diff & (1<= 0; k--){ + if(up[a][k] != up[b][k]){ + a = up[a][k]; + b = up[b][k]; + } + } + return up[a][0]; + } + + int stepUp(int u, int k) { + for (int i=16; i>=0; i--) { + if ((k>>i)&1) { + u = up[u][i]; + } + } + return u; + } + + ll dist(int a, int b) + { + int c = lca(a,b); + return distRoot[a] + distRoot[b] - 2*distRoot[c]; + } + + vector findMedian(int n, vector>& edges, vector>& queries) { + for (int i = 0; i < n-1; i++) + { + int u = edges[i][0], v = edges[i][1], w = edges[i][2]; + adj[u].push_back({v,w}); + adj[v].push_back({u,w}); + } + + depth[0] = 0; + distRoot[0] = 0; + dfs(0, 0); + + for(int k = 1; k <= LOGN; k++) { + for(int v = 0; v < n; v++) { + up[v][k] = up[up[v][k-1]][k-1]; + } + } + + vectorrets; + for (auto& q: queries) + { + int u = q[0], v = q[1]; + int c = lca(u,v); + ll total = dist(u,c)+dist(c,v); + + int step1 = depth[u]-depth[c]; + int step2 = depth[v]-depth[c]; + + int low = 0, high = step1+step2; + int k; + while (low < high) { + int mid = low + (high-low)/2; + ll d; + if (mid <= step1) { + k = stepUp(u, mid); + d = distRoot[u] - distRoot[k]; + } else { + k = stepUp(v, step2 - (mid-step1)); + d = total - (distRoot[v] - distRoot[k]); + } + if (d >= total*0.5) + high = mid; + else + low = mid+1; + } + int step = low; + if (step<=step1) + rets.push_back(stepUp(u, step)); + else + rets.push_back(stepUp(v, step2-(step-step1))); + } + + return rets; + } +}; diff --git a/Binary_Search/3585.Find-Weighted-Median-Node-in-Tree/Readme.md b/Binary_Search/3585.Find-Weighted-Median-Node-in-Tree/Readme.md new file mode 100644 index 000000000..87b7b21b1 --- /dev/null +++ b/Binary_Search/3585.Find-Weighted-Median-Node-in-Tree/Readme.md @@ -0,0 +1,7 @@ +### 3585.Find-Weighted-Median-Node-in-Tree + +对于任何一个query,我们只需要找到u到v路径(途中经过LCA的点记作c),假设路径的总步长是d,路径的总权重和是total。我们只需要在[0,d]之间进行二分搜索一个合适的步数k:即从u走k步,恰好走过的路径长度超过total的一半。 + +注意,我们在二分搜索对k进行判定的时候,需要分类讨论k是否在u到c的路径上,还是c到v的路径上。即看是否`dist(u,c) >= total * 0.5`. 如果k是在u到c的路径上,那么经过的路径长度就是dist(u,k)。如果k是在c到v的路径上,那么经过的路径长度就是dist(u,c)+dist(c,k)。 + +根据binary lifting的算法,树里任意两个节点之间的距离都可以用log(n)的时间求解。 diff --git a/Readme.md b/Readme.md index 0e1dda6f8..21b6d522d 100644 --- a/Readme.md +++ b/Readme.md @@ -108,6 +108,7 @@ [2846.Minimum-Edge-Weight-Equilibrium-Queries-in-a-Tree](https://github.com/wisdompeak/LeetCode/tree/master/Binary_Search/2846.Minimum-Edge-Weight-Equilibrium-Queries-in-a-Tree) (H) [2851.String-Transformation](https://github.com/wisdompeak/LeetCode/tree/master/Dynamic_Programming/2851.String-Transformation) (H+) [3553.Minimum-Weighted-Subgraph-With-the-Required-Paths-II](https://github.com/wisdompeak/LeetCode/tree/master/Binary_Search/3553.Minimum-Weighted-Subgraph-With-the-Required-Paths-II) (H) +[3585.Find-Weighted-Median-Node-in-Tree](https://github.com/wisdompeak/LeetCode/tree/master/Binary_Search/3585.Find-Weighted-Median-Node-in-Tree) (H) * ``Binary Search by Value`` [410.Split-Array-Largest-Sum](https://github.com/wisdompeak/LeetCode/tree/master/Dynamic_Programming/410.Split-Array-Largest-Sum) (H-) [774.Minimize-Max-Distance-to-Gas-Station](https://github.com/wisdompeak/LeetCode/tree/master/Priority_Queue/774.Minimize-Max-Distance-to-Gas-Station) (H)