class FenwickTree:
def __init__(self, n: int):
self.tree = [0] * (n + 1) # 使用下标 1 到 n
# a[i] 增加 val
# 1 <= i <= n
def update(self, i: int, val: int) -> None:
while i < len(self.tree):
self.tree[i] += val
i += i & -i
# 计算前缀和 a[1] + ... + a[i]
# 1 <= i <= n
def pre(self, i: int) -> int:
res = 0
while i > 0:
res += self.tree[i]
i &= i - 1
return res
class Solution:
def treeQueries(self, n: int, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
g = [[] for _ in range(n + 1)]
for x, y, _ in edges:
g[x].append(y)
g[y].append(x)
in_ = [0] * (n + 1)
out = [0] * (n + 1)
clock = 0
def dfs(x: int, fa: int) -> None:
nonlocal clock
clock += 1
in_[x] = clock # 进来的时间
for y in g[x]:
if y != fa:
dfs(y, x)
out[x] = clock # 离开的时间
dfs(1, 0)
weight = [0] * (n + 1)
diff = FenwickTree(n)
def update(x: int, y: int, w: int) -> None:
# 保证 y 是 x 的儿子
if in_[x] > in_[y]:
y = x
d = w - weight[y] # 边权的增量
weight[y] = w
# 把子树 y 中的最短路长度都增加 d（用差分树状数组维护）
diff.update(in_[y], d)
diff.update(out[y] + 1, -d)
for e in edges:
update(*e)
ans = []
for q in queries:
if q[0] == 1:
update(*q[1:])
else:
ans.append(diff.pre(in_[q[1]]))
return ans
class FenwickTree {
private final int[] tree;
public FenwickTree(int n) {
tree = new int[n + 1]; // 使用下标 1 到 n
}
// a[i] 增加 val
// 1 <= i <= n
public void update(int i, int val) {
for (; i < tree.length; i += i & -i) {
tree[i] += val;
}
}
// 求前缀和 a[1] + ... + a[i]
// 1 <= i <= n
public int pre(int i) {
int res = 0;
for (; i > 0; i &= i - 1) {
res += tree[i];
}
return res;
}
}
class Solution {
public int[] treeQueries(int n, int[][] edges, int[][] queries) {
List<Integer>[] g = new ArrayList[n + 1];
Arrays.setAll(g, i -> new ArrayList<>());
for (int[] e : edges) {
int x = e[0];
int y = e[1];
g[x].add(y);
g[y].add(x);
}
int[] in = new int[n + 1];
int[] out = new int[n + 1];
dfs(1, 0, g, in, out);
int[] weight = new int[n + 1];
FenwickTree diff = new FenwickTree(n);
for (int[] e : edges) {
update(e[0], e[1], e[2], in, out, weight, diff);
}
List<Integer> ans = new ArrayList<>();
for (int[] q : queries) {
if (q[0] == 1) {
update(q[1], q[2], q[3], in, out, weight, diff);
} else {
ans.add(diff.pre(in[q[1]]));
}
}
return ans.stream().mapToInt(i -> i).toArray();
}
private int clock = 0;
private void dfs(int x, int fa, List<Integer>[] g, int[] in, int[] out) {
in[x] = ++clock; // 进来的时间
for (int y : g[x]) {
if (y != fa) {
dfs(y, x, g, in, out);
}
}
out[x] = clock; // 离开的时间
}
private void update(int x, int y, int w, int[] in, int[] out, int[] weight, FenwickTree diff) {
// 保证 y 是 x 的儿子
if (in[x] > in[y]) {
y = x;
}
int d = w - weight[y]; // 边权的增量
weight[y] = w;
// 把子树 y 中的最短路长度都增加 d（用差分树状数组维护）
diff.update(in[y], d);
diff.update(out[y] + 1, -d);
}
}
