以下の bonus 制約でも解けるので、その解説を書きます。

• bonus: \(N \le 2\times 10^5\)
• bonus: 辺 \((i, P_i)\) が一般の functional graph を成す
  • 元問題における \(P_i \neq i\) と \(P\) が相異なる制約がないのに相当

関連する問題として、まず以下の問題を考えます。

配列 \(a = (a_0, a_1, \dots, a_{n-1})\) が与えられる。 次のクエリに高速に答えよ。

• 1 \(i\) \(x\)
  • \(a_i \gets x\) で更新
• 2 \(l\)
  • \(\max_{l < r < n} \sum_{i=l}^{r-1} a_i\) を出力
  • 要するに、\(l\) から始まる（空でない）区間の和のうちの最大値

これは、以下のような情報を管理することでモノイドを作れるため、セグ木で解くことができます。

• 区間 \([l, r)\) の和 sum
• \(l < i \le r\) における区間 \([l, i)\) の和の最大値 max

これらのモノイド積を考えます。

(left.sum, left.max) と (right.sum, right.max) の積は (left.sum + right.sum, max(left.max, left.sum + right.max)) で得られます。

単位元は \((0, -\infty)\) です。

note: 空区間を許すなら \((0, 0)\) となりそうです。また、適切に管理することで、prefix sum だけでなく subarray sum の最大値なども計算できます。

さて、元の問題について考えます。

functional graph の各辺について上記のモノイドの値を乗せます。 具体的には、頂点 \(v\) から \(P_v\) への辺には \((C_{P_v}, C_{P_v})\) を乗せます。

与えられた頂点 \(v\) から \(K\) 回辿る間に見た辺にある値たちのモノイド積を計算することで、頂点 \(v\) から始めた際のスコアの最大値が得られます。 この値は、ダブリングを用いることで \(\langle O(N\log(K)), O(\log(K))\rangle\) で処理できるため、各 \(v\) について行うことで全体の答えが \(O(N\log(K))\) 時間で求まります。

• 実装例：提出 #34409729, 21 ms
  • ここではモノイド積ではなく半群積のような実装になっていますが、本質は同じです。