Hard
题目描述
给定一个整数 n 和一个根节点为 0 的无向加权树，该树有 n 个节点，编号从 0 到 n - 1。树由长度为 n - 1 的二维数组 edges 表示，其中 edges[i] = [ui, vi, wi] 表示节点 ui 到 vi 之间有一条权重为 wi 的边。
加权中位数节点定义为：在从 ui 到 vi 的路径上的第一个节点 x，使得从 ui 到 x 的边权重总和大于或等于总路径权重的一半。
给定一个二维整数数组 queries。对于每个 queries[j] = [uj, vj]，确定从 uj 到 vj 路径上的加权中位数节点。
返回一个数组 ans，其中 ans[j] 是 queries[j] 的加权中位数的节点索引。
示例 1：
输入：n = 2, edges = [[0,1,7]], queries = [[1,0],[0,1]]
输出：[0,1]
解释：
- 查询 [1,0]：路径 1→0，边权重 [7]，总权重 7，一半为 3.5。从 1 到 0 的权重和为 7 >= 3.5，中位数是节点 0。
- 查询 [0,1]：路径 0→1，边权重 [7]，总权重 7，一半为 3.5。从 0 到 1 的权重和为 7 >= 3.5，中位数是节点 1。
示例 2：
输入：n = 3, edges = [[0,1,2],[2,0,4]], queries = [[0,1],[2,0],[1,2]]
输出：[1,0,2]
示例 3：
输入：n = 5, edges = [[0,1,2],[0,2,5],[1,3,1],[2,4,3]], queries = [[3,4],[1,2]]
输出：[2,2]
约束条件：
- 2 <= n <= 10^5
- edges.length == n - 1
- edges[i] == [ui, vi, wi]
- 0 <= ui, vi < n
- 1 <= wi <= 10^9
- 1 <= queries.length <= 10^5
- queries[j] == [uj, vj]
- 0 <= uj, vj < n
- 输入保证 edges 表示一个有效的树
解题思路
这是一道综合性很强的树上路径查询问题。我们需要对每个查询找到路径上的加权中位数节点。
核心思路：
预处理阶段：构建树的邻接表，并使用二进制提升（Binary Lifting）技术预处理每个节点的祖先信息和到祖先的累积权重。
LCA查找：对于每个查询 [u,v]，首先找到它们的最近公共祖先（LCA）。
路径分析：从 u 到 v 的路径可以分解为 u→LCA 和 LCA→v 两段。计算总路径权重，目标是找到第一个使累积权重 ≥ 总权重/2 的节点。
二分查找定位：
- 首先检查中位数是否在 u→LCA 路径上
- 如果不在，则在 v→LCA 路径上查找
- 使用二进制提升快速跳跃，在 O(log n) 时间内定位
算法步骤：
- 使用 DFS 预处理每个节点的深度、父节点和边权
- 构建二进制提升表，存储 2^k 级祖先和对应的累积权重
- 对每个查询：计算 LCA，确定总路径权重，然后二分查找中位数节点
时间复杂度优化：通过二进制提升避免了朴素的逐节点遍历，将单次查询的时间复杂度从 O(n) 优化到 O(log n)。
代码实现
class Solution {
private:
vector<vector<pair<int, long long>>> adj;
vector<vector<int>> up;
vector<vector<long long>> upWeight;
vector<int> depth;
int LOG;
void dfs(int u, int p, int d) {
depth[u] = d;
up[u][0] = p;
for (int i = 1; i < LOG; i++) {
if (up[u][i-1] != -1) {
up[u][i] = up[up[u][i-1]][i-1];
upWeight[u][i] = upWeight[u][i-1] + upWeight[up[u][i-1]][i-1];
}
}
for (auto [v, w] : adj[u]) {
if (v != p) {
upWeight[v][0] = w;
dfs(v, u, d + 1);
}
}
}
int lca(int u, int v) {
if (depth[u] < depth[v]) swap(u, v);
int diff = depth[u] - depth[v];
for (int i = 0; i < LOG; i++) {
if ((diff >> i) & 1) {
u = up[u][i];
}
}
if (u == v) return u;
for (int i = LOG - 1; i >= 0; i--) {
if (up[u][i] != up[v][i]) {
u = up[u][i];
v = up[v][i];
}
}
return up[u][0];
}
long long getPathWeight(int u, int ancestor) {
if (u == ancestor) return 0;
long long total = 0;
int diff = depth[u] - depth[ancestor];
for (int i = 0; i < LOG; i++) {
if ((diff >> i) & 1) {
total += upWeight[u][i];
u = up[u][i];
}
}
return total;
}
int findMedianOnPath(int start, int end, long long target) {
if (start == end) return start;
long long sum = 0;
int curr = start;
while (curr != end) {
int diff = depth[curr] - depth[end];
int jump = 0;
// Find largest jump that doesn't exceed target
for (int i = LOG - 1; i >= 0; i--) {
if ((diff >> i) & 1) {
int next = curr;
long long nextSum = sum;
for (int j = i; j >= 0; j--) {
if ((diff >> j) & 1) {
nextSum += upWeight[next][j];
next = up[next][j];
}
}
if (nextSum >= target) {
// Binary search within this range
for (int j = i; j >= 0; j--) {
if (((diff >> j) & 1) && sum + upWeight[curr][j] < target) {
sum += upWeight[curr][j];
curr = up[curr][j];
diff -= (1 << j);
}
}
// Move one more step
if (curr != end) {
sum += upWeight[curr][0];
curr = up[curr][0];
}
return curr;
}
break;
}
}
sum += upWeight[curr][0];
curr = up[curr][0];
}
return curr;
}
public:
vector<int> findMedian(int n, vector<vector<int>>& edges, vector<vector<int>>& queries) {
LOG = 20;
adj.resize(n);
up.assign(n, vector<int>(LOG, -1));
upWeight.assign(n, vector<long long>(LOG, 0));
depth.resize(n);
for (auto& edge : edges) {
int u = edge[0], v = edge[1];
long long w = edge[2];
adj[u].push_back({v, w});
adj[v].push_back({u, w});
}
dfs(0, -1, 0);
vector<int> result;
for (auto& query : queries) {
int u = query[0], v = query[1];
int l = lca(u, v);
long long pathWeight = getPathWeight(u, l) + getPathWeight(v, l);
long long target = (pathWeight + 1) / 2; // Ceiling division
long long uToLCA = getPathWeight(u, l);
if (uToLCA >= target) {
result.push_back(findMedianOnPath(u, l, target));
} else {
long long remaining = target - uToLCA;
long long vToLCA = getPathWeight(v, l);
result.push_back(findMedianOnPath(v, l, vToLCA - remaining + 1));
}
}
return result;
}
};
class Solution:
def findMedian(self, n: int, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
from collections import defaultdict
LOG = 20
adj = defaultdict(list)
up = [[-1] * LOG for _ in range(n)]
up_weight = [[0] * LOG for _ in range(n)]
depth = [0] * n
for u, v, w in edges:
adj[u].append((v, w))
adj[v].append((u, w))
def dfs(u, p, d):
depth[u] = d
up[u][0] = p
for i in range(1, LOG):
if up[u][i-1] != -1:
up[u][i] = up[up[u][i-1]][i-1]
up_weight[u][i] = up_weight[u][i-1] + up_weight[up[u][i-1]][i-1]
for v, w in adj[u]:
if v != p:
up_weight[v][0] = w
dfs(v, u, d + 1)
def lca(u, v):
if depth[u] < depth[v]:
u, v = v, u
diff = depth[u] - depth[v]
for i in range(LOG):
if (diff >> i) & 1:
u = up[u][i]
if u == v:
return u
for i in range(LOG-1, -1, -1):
if up[u][i] != up[v][i]:
u = up[u][i]
v = up[v][i]
return up[u][0]
def get_path_weight(u, ancestor):
if u == ancestor:
return 0
total = 0
diff = depth[u] - depth[ancestor]
for i in range(LOG):
if (diff >> i) & 1:
total += up_weight[u][i]
u = up[u][i]
return total
def find_median_on_path(start, end, target):
if start == end:
return start
curr = start
total = 0
while curr != end:
if total + up_weight[curr][0] >= target:
return up[curr][0]
total += up_weight[curr][0]
curr = up[curr][0]
return curr
dfs(0, -1, 0)
result = []
for u, v in queries:
l = lca(u, v)
path_weight = get_path_weight(u, l) + get_path_weight(v, l)
target = (path_weight + 1) // 2
u_to_lca = get_path_weight(u, l)
if u_to_lca >= target:
result.append(find_median_on_path(u, l, target))
else:
remaining = target - u_to_lca
v_to_lca = get_path_weight(v, l)
result.append(find_median_on_path(v, l, v_to_lca - remaining + 1))
return result
public class Solution {
private List<(int, long)>[] adj;
private int[,] up;
private long[,] upWeight;
private int[] depth;
private int LOG = 20;
public int[] FindMedian(int n, int[][] edges, int[][] queries) {
adj = new List<(int, long)>[n];
up = new int[n, LOG];
upWeight = new long[n, LOG];
depth = new int[n];
for (int i = 0; i < n; i++) {
adj[i] = new List<(int, long)>();
for (int j = 0; j < LOG; j++) {
up[i, j] = -1;
}
}
foreach (var edge in edges) {
int u = edge[0], v = edge[1];
long w = edge[2];
adj[u].Add((v, w));
adj[v].Add((u, w));
}
DFS(0, -1, 0);
List<int> result = new List<int>();
foreach (var query in queries) {
int u = query[0], v = query[1];
int l = LCA(u, v);
long pathWeight = GetPathWeight(u, l) + GetPathWeight(v, l);
long target = (pathWeight + 1) / 2;
long uToLCA = GetPathWeight(u, l);
if (uToLCA >= target) {
result.Add(FindMedianOnPath(u, l, target));
} else {
long remaining = target - uToLCA;
long vToLCA = GetPathWeight(v, l);
result.Add(FindMedianOnPath(v, l, vToLCA - remaining + 1));
}
}
return result.ToArray();
}
private void DFS(int u, int p, int d) {
depth[u] = d;
up[u, 0] = p;
for (int i = 1; i < LOG; i++) {
if (up[u, i-1] != -1) {
up[u, i] = up[up[u, i-1], i-1];
upWeight[u, i] = upWeight[u, i-1] + upWeight[up[u, i-1], i-1];
}
}
foreach (var (v, w) in adj[u]) {
if (v != p) {
upWeight[v, 0] = w;
DFS(v, u, d + 1);
}
}
}
private int LCA(int u, int v) {
if (depth[u] < depth[v]) {
int temp = u; u = v; v = temp;
}
int diff = depth[u] - depth[v];
for (int i = 0; i < LOG; i++) {
if ((diff >> i & 1) == 1) {
u = up[u, i];
}
}
if (u == v) return u;
for (int i = LOG - 1; i >= 0; i--) {
if (up[u, i] != up[v, i]) {
u = up[u, i];
v = up[v, i];
}
}
return up[u, 0];
}
private long GetPathWeight(int u, int ancestor) {
if (u == ancestor) return 0;
long total = 0;
int diff = depth[u] - depth[ancestor];
for (int i = 0; i < LOG; i++) {
if ((diff >> i & 1) == 1) {
total += upWeight[u, i];
u = up[u, i];
}
}
return total;
}
private int FindMedianOnPath(int start, int end, long target) {
if (start == end) return start;
int curr = start;
long total = 0;
while (curr != end) {
if (total + upWeight[curr, 0] >= target) {
return up[curr, 0];
}
total += upWeight[curr, 0];
curr = up[curr, 0];
}
return curr;
}
}
var findMedian = function(n, edges, queries) {
const graph = Array(n).fill(null).map(() => []);
for (const [u, v, w] of edges) {
graph[u].push([v, w]);
graph[v].push([u, w]);
}
function findPath(start, end) {
const visited = new Set();
const path = [];
function dfs(node, target, currentPath) {
if (visited.has(node)) return false;
visited.add(node);
currentPath.push(node);
if (node === target) {
path.push(...currentPath);
return true;
}
for (const [neighbor, weight] of graph[node]) {
if (dfs(neighbor, target, currentPath)) {
return true;
}
}
currentPath.pop();
return false;
}
dfs(start, end, []);
return path;
}
function getPathWeights(path) {
const weights = [];
for (let i = 0; i < path.length - 1; i++) {
const curr = path[i];
const next = path[i + 1];
for (const [neighbor, weight] of graph[curr]) {
if (neighbor === next) {
weights.push(weight);
break;
}
}
}
return weights;
}
const result = [];
for (const [u, v] of queries) {
const path = findPath(u, v);
const weights = getPathWeights(path);
const totalWeight = weights.reduce((sum, w) => sum + w, 0);
const half = totalWeight / 2;
let currentSum = 0;
let medianNode = path[0];
for (let i = 0; i < weights.length; i++) {
currentSum += weights[i];
if (currentSum >= half) {
medianNode = path[i + 1];
break;
}
}
result.push(medianNode);
}
return result;
};
复杂度分析
| 复杂度类型 | 分析 |
|---|---|
| 时间复杂度 | O((n + q) log n)，其中 n 是节点数，q 是查询数。预处理需要 O(n log n)，每次查询需要 O(log n) |
| 空间复杂度 | O(n log n)，用于存储二进制提升表和累积权重信息 |