ヒント → https://atcoder.jp/contests/agc063/editorial/6736

---

以下では「操作方法」といえば，各頂点を選ぶ回数の定め方のことを指すものとし，操作手順の違いは区別しないこととします．また，「部分木 \(v\) 内での操作」といえば，\(i\) と \(P_i\) がともに部分木 \(v\) に含まれるような操作を指すこととします．

また，入力で与えられる \(A_i\) の初期値を \(a_i\) と書くことにします．

---

[1] 木 dp と多項式による表示

本問題は，計算量を無視すれば，自然な木 dp で解けます．まずはこの計算手順を確認します．計算手順を分解して次のようにいくつかの dp テーブルを定義します．

• \(\mathrm{dp}_1[v][k]\)：部分木 \(v\) 内での操作方法であって，\(A_v\) がちょうど \(k\) 回 \(+r\) の対象となるようなものの個数．
• \(\mathrm{dp}_2[v][k]\)：部分木 \(v\) 内での操作方法であって，\(A_v\) がちょうど \(+k\) されるようなものの個数．
• \(\mathrm{dp}_3[v][k]\)：部分木 \(v\) 内での操作方法であって，\(A_v=k\) になるようなものの個数．
• \(\mathrm{dp}_4[v][k]\)：部分木 \(v\) 内および，頂点 \(v\) から \(P_v\) への操作方法であって，\(A_v\) がちょうど \(k\) 回 \(-1\) の対象となるようなものの個数．

これらの計算方法は以下の通りです．\(v\) の子全体を \(\mathrm{ch}(v)\) と書くことにします．

• \(\mathrm{dp}_1[v]\) は，すべての \(w\in \mathrm{ch}(v)\) に対する \(\mathrm{dp}_4[w]\) の畳み込み．
• \(\mathrm{dp}_2[v][rk] = \mathrm{dp}_1[v][k]\)．
• \(\mathrm{dp}_3[v][a_v+k] = \mathrm{dp}_2[v][k]\)．
• \(\mathrm{dp}_4[v][k] = \sum_{i\geq k}\mathrm{dp}_3[v][i]\)．

畳み込み操作が出てくるので，配列の代わりに多項式を考えましょう．\(F_{v,1}(x) = \sum_{k}\mathrm{dp}[v][k]x^k\) とします．同様に \(F_{v,2}(x)\), \(F_{v,3}(x)\), \(F_{v,4}(x)\) も定めます．すると，上の計算は次のように表すことができます．

• \(F_{v,1}(x) = \prod_{w\in\mathrm{ch}(v)}F_{w,4}(x)\)．
• \(F_{v,2}(x) = F_{v,1}(x^r)\)．
• \(F_{v,3}(x) = x^{a_v}F_{v,2}(x)\)．
• \(F_{v,4}(x) = F_{v,3}(x)+\dfrac{F_{v,3}(1)-F_{v,3}(x)}{1-x}\)．

\(F_{v,3}\) から \(F_{v,4}\) を得る操作のみ少し非自明ですが，\(x^k\) を \(1+\cdots +x^k\) に置き換える線形な操作がこのように書けます．

これを葉側から順に計算して，\(v\) が根のときの \(F_{v,3}(1)\) を答えればよいことになります．

---

[2] 変数変換の利用

[1] の解法をそのまま利用しようとすると，答の計算のために膨大な係数列を保持する必要が出てきてしまい，高速に解くことができません．

特に答を得たり \(F_{v,4}\) を計算するために，\(x=1\) の代入を考えているのが厄介で，このため多項式の高次の項も無視できなくなっています．

そこで，\(G(x)=F(1+x)\) というタイプの変数変換を考えます．つまり，\(G_{v,1}(x) = F_{v,1}(1+x)\) として \(G_{v,1}\) を定義し，同様に \(G_{v,2}, G_{v,3}, G_{v,4}\) を定義します．

\(F_{v,i}\) の計算手順を適切な式変形により書き直せば，\(G_{v,i}\) は次の手順で計算できることが分かります．

• \(G_{v,1}(x) = \prod_{w\in\mathrm{ch}(v)}G_{w,4}(x)\)．
• \(G_{v,2}(x) = G_{v,1}((1+x)^r-1)\)．
• \(G_{v,3}(x) = (1+x)^{a_v}G_{v,2}(x)\)．
• \(G_{v,4}(x) = G_{v,3}(x)+\dfrac{G_{v,3}(0)-G_{v,3}(x)}{-x}\)．

求めるべきは \(v\) が根のときの \(G_{v,3}(0)\)，つまり，\(G_{v,3}(x)\pmod{x^1}\) です．

\(n\) を正整数とするとき次が成り立つことが容易に分かります：

• \(G_{v,1}(x) \pmod{x^n}\) は \(w\in \mathrm{ch}(v)\) に対する \(G_{w,4}(x) \pmod{x^n}\) から定まる．
• \(G_{v,2}(x) \pmod{x^n}\) は \(G_{1,v}(x)\pmod{x^n}\) から定まる．
• \(G_{v,3}(x)\pmod{x^n}\) は \(G_{v,2}(x)\pmod{x^n}\) から定まる．
• \(G_{v,4}(x) \pmod{x^n}\) は \(G_{v,3}(x) \pmod{x^{n+1}}\) から定まる．

\(G_{v,3}\) から \(G_{v,4}\) を得る部分のみ指数がずれていることに注意してください．

したがって，深さ \(d\) の頂点 \(v\) において，

• \(G_{v,1} \pmod{x^{1+d}}\)
• \(G_{v,2} \pmod{x^{1+d}}\)
• \(G_{v,3} \pmod{x^{1+d}}\)
• \(G_{v,4} \pmod{x^{d}}\)

を求めることを考えると，これは葉側から順に計算できます．

各頂点で \(O(N)\) 個の係数しか扱わなくて良くなったため，[1] で述べた解法よりも高速な計算が可能になります．各ステップを適切に実装すれば，本問題を \(O(N^3)\) 時間で解くことができます．

\(G((1+x)^r-1)\) の計算のみ少し難しいですが，\(H(x) = G(x-1)\) を計算してから \(H((1+x)^r)\) を求めるという手順をとり，適切に二項定理を使えばこの合成は \(O(d^2)\) 時間で計算できます．

---

[3] 参考

• \(G(x) = F(x+1)\) を考える代わりに，高階微分係数の列 \(\bigl(F(1),F'(1),F''(1),\ldots\bigr)\) を考えても全く同じことになります．
• 多項式に関する高度なアルゴリズムを用いれば，\(O(N^2\log^2 N)\) 時間で解くこともできます．最も難しい部分は合成 \(G((1+x)^r-1)\) の計算だと思います．これも polynomial taylor shift, \(\exp\) や \(\log\) との合成などに帰着すれば，\(O(d\log^2d)\) 時間で行えます．参考リンク