- Define precisely what one DP state represents.
- Establish the base cases before transitions are evaluated.
- Process states in dependency order and combine only already-known values.
Code notes
- 58 lines of C++ from the credited upstream file yosupo-tree-path-composite-sum.test.cpp.
- The implementation visibly relies on sequence storage, ordered lookup, cached states.
- 1 loop block detected.
Complexity
Multiply the number of reachable states by the work performed for each transition, then include the stored state table in memory usage.
Check the problem constraints before deciding whether this complexity will pass.
Use this to learn the idea, then write your own version.
1#define PROBLEM "https://judge.yosupo.jp/problem/tree_path_composite_sum"23#include "../../template/template.hpp"45#include "../../graph/graph-template.hpp"67#include "../../tree/rerooting.hpp"89#include "../../modint/montgomery-modint.hpp"10#include "../../modulo/binomial.hpp"1112using namespace Nyaan;13using mint = LazyMontgomeryModInt<998244353>;1415using vm = vector<mint>;16using vvm = vector<vm>;17Binomial<mint> C;18 19using namespace Nyaan;20 21void q() {22 inl(N);23 vm A(N);24 in(A);25 26 using pm = pair<mint, mint>;27 vvi g(N);28 map<pi, pm> mp;29 rep(_, N - 1) {30 ini(u, v, b, c);31 g[u].push_back(v);32 g[v].push_back(u);33 mp[minmax(u, v)] = {b, c};34 }35 36 37 using T = pair<mint, int>;38 39 T leaf = {0, 0};40 41 auto f1 = [&](T x, T y) -> T { return {x.fi + y.fi, x.se + y.se}; };42 43 auto f2 = [&](T xn, int c, int p) -> T {44 auto [x, n] = xn;45 auto [a, b] = mp[minmax(c, p)];46 return make_pair((x + A[c]) * a + b * (n + 1), n + 1);47 };48 Rerooting<T, decltype(g), decltype(f1), decltype(f2)> dp(g, f1, f2, leaf);49 auto ans = dp.dp;50 rep(i, N) cout << (ans[i].fi + A[i]) << " \n"[i + 1 == N];51}52 53void Nyaan::solve() {54 int t = 1;55 56 while (t--) q();57}58