Tổng tiền tố trên cây

Hướng tiếp cận

Với mỗi đỉnh $u$, gọi $\text{pref}[u]$ là tổng trọng số trên đường đi từ gốc đến $u$. Khi đó:

$$\text{pref}[u]=\text{pref}[\text{parent}(u)]+w_u$$

với $\text{pref}[1]=w_1$. Sau khi tính sẵn $\text{pref}[u]$ cho mọi đỉnh bằng một lần DFS, mỗi truy vấn chỉ cần in ra $\text{pref}[u]$ trong $O(1)$.

Nhận xét quan trọng

1. Vì $|w_u|\le {10}^9$ và cây có thể sâu tới $2·{10}^5$ đỉnh, tổng có thể đạt đến $2·{10}^{14}$ — vượt phạm vi int. Phải dùng long long.
2. Cây có thể là một đường thẳng dài (n = 2·10⁵). DFS đệ quy có thể gây tràn stack — nên dùng DFS lặp (stack thủ công) hoặc tăng sys.setrecursionlimit trong Python (nhưng đệ quy Python sâu 2·10⁵ vẫn rất chậm và có thể vẫn lỗi). Triển khai DFS lặp là an toàn nhất.

Thuật toán

1. Đọc $n$, mảng trọng số $w$, các cạnh, dựng danh sách kề.
2. Chạy DFS lặp từ đỉnh $1$:
   • Dùng một stack chứa $(u,\text{parent})$.
   • Khi pop một đỉnh $u$, gán $\text{pref}[u]=\text{pref}[\text{parent}]+w_u$ (riêng gốc thì $\text{pref}[1]=w_1$).
   • Đẩy các con của $u$ (mọi hàng xóm khác parent) vào stack.
3. Với mỗi truy vấn $u$, in ra $\text{pref}[u]$.

Độ phức tạp

• Thời gian: $O(n+q)$.
• Bộ nhớ: $O(n)$.