はじめに
問題リンク
「How much did kotamanegi earn?」writerのvwxyzです。
ステップが多く、グラフの理解、数論の考察、丁寧な場合分けが要求され、難しい問題だったと思います。
コンテスト開始後90分間くらい問題文のミスに気付かず、全体順位にも少し影響が出てしまいました。申し訳ありません。
解説
各社員$i$（$1\leq i\leq N$）について、ミスを起こすと社員$i$にも責任が生じるようなすべての社員（社員$i$本人も含まれる）を社員$i$の部下と呼ぶことにします。$i=1,2,\cdots ,N$に対して$E_i$に$N^K$をかけたものを$S_i$と表すことにすると、$S_i$は社員$i$が罰金を支払うようなすべての場合についての罰金の総和になります。
ゼータ変換によるminの畳み込みの要領で
$i=1,2,\cdots ,N$に対して
$\zeta (A_i)=\displaystyle\sum_{iの部下j}A_j$
$\zeta (S_i)=\displaystyle\sum_{iの部下j}S_j$
と表すことにすると、$\zeta (A_i)$、$\zeta (S_i)$は$O(N)$で求められます。
罰金額の定め方から
$\zeta (A_i)^K=\zeta (S_i)$（・・・①）
の関係が成り立つことがわかります。
離散対数問題を解く
$P=10^9+7$、$p=500000003$（$P=2\times p+1$）とおくと、$P$も$p$もどちらも素数です。
まずは$K\neq p$の場合を考えます。このとき、
$P-1$を素因数分解すると$2\times p$
$K$は奇数
$1\leq K\leq 10^9<2p$
であることから、$K$と$P-1$は互いに素であることがわかるので、$modP-1$における$K$の逆元を$K_{inv}$とおくことができます。①の両辺を$K_{inv}$乗することで
$\zeta (A_i)\equiv\zeta (S_i)^{K_{inv}}\quad(modP)$（・・・②）
がわかり、また、②の両辺を$K$乗すると①になるので、①と②は同値です。したがって$\zeta (A_i)$を一意に定めることができ、そこから$A_i$も$O(N)$で計算することができます。
コーナーケースの場合分け
次に$K=p$の場合を考えます。$x=0,1,\cdots,P-1$に対して$x^K$が$modP$でどのような値を取りうるかについて考えます。
$x=0$なら$x^K$は$0$になります。
$x\neq 0$ならフェルマーの小定理より
$x^{2K}\equiv1(modP)$
なので
$(x^K-1)(x^K+1)\equiv0(modP)$
となり、$x^K$は$modP$において$1$か$P-1$のどちらかに等しく、また$K$は素数であることから
$x^K\equiv-(-x)^K\equiv-(P-x)^K(modP)$（・・・③）
なので、$x^K$は$x=1,2,\cdots,P-1$のうち$p$個が$1$、もう$p$個が$P-1$となる（・・・④）ことがわかります。したがって、$\zeta (S_i)$に$0$、$1$、$P-1$以外の値があれば入手した情報に誤りがあり、$0$、$1$、$P-1$のみからなる場合は条件を満たすような$\zeta (A_i)$が存在し、$A_i$も存在することがわかります。kotamanegi社長の年収は
$\zeta (A_1)-\displaystyle\sum_{直属の上司が社長であるような社員i}\zeta (A_i)(modP)$（・・・⑤）
なので、これの最小値を考えます。直属の上司が社長であるような社員iのうち$\zeta (A_i)\neq0$であるような（$\zeta (A_i)\equiv 0(modP)$なら⑤の式に影響を及ぼさないので考える必要がないため）社員を「右腕社員」と呼ぶことにします。
右腕社員$i$、$j$で$\zeta (A_i)^K\equiv 1(modP)$、$\zeta (A_j)^K\equiv P-1(modP)$となるような二人がいたとします。$1\leq a\leq P-1$を満たす整数$a$で$\zeta (A_i)+\zeta (A_j)\equiv a(modP)$を満たすように$\zeta (A_i)$と$\zeta (A_j)$を設定できないと仮定します。$X_1=\lbrace x\mid 0\leq x\leq P-1\land x^K\equiv 1(modP)\rbrace$、$X_2=\lbrace x\mid 0\leq x\leq P-1\land x^K\equiv 1(modP)\rbrace$とおくと、③と④から
$X_1=\lbrace a-x(modP)\mid x\in X_2\rbrace$
$X_1=\lbrace -x(modP)\mid x\in X_2\rbrace$
となるので$X_1$からある要素$x$を取ると$x+a$も$X_1$に含まれることがわかり、$x$と$P$は互いに素であることから$X_1=\lbrace x\mid 0\leq x\leq P-1\rbrace$となり、矛盾します。
また、$1 \in X_1$、$P-1\in X_2$であり、$1+(P-1)\equiv 0(modP)$なので、$\zeta (A_i)^K\equiv 1 (modP)$、$\zeta (A_j)^K\equiv P-1(modP)$となるような二人の右腕社員$i$、$j$がいれば$\zeta (A_1)$の値によらず⑤を$0$にできることがわかります。
同様に考察することで、$\zeta (A_i)^K\equiv \zeta (A_j)^K(modP)$となるような二人の右腕社員$i$、$j$がいれば$\zeta (A_1)\neq 0$のときには⑤を$0$にできることがわかります。
（実際には例えば$\zeta (A_1)$が$0$、$1$、$P-1$の場合についてのみ証明すれば十分です。）
以上のことから、右腕社員が$3$人以上いる場合は答えが$0$になることがわかります。
右腕社員が$2$人いる場合は$\zeta (A_i)^K\equiv \zeta (A_j)^K(modP)\land\zeta (A_1)\equiv 0(modP)$の場合は$1$、そうでない場合は$0$になることがわかります。
右腕社員が$1$人の場合も同様に答えを絞り込んでいくことができます。
$x=1,2,\cdots,10$と$x=P-10,P-9,\cdots,P-1$の場合に$x^K(modP)$を計算すれば答えを求めるのには十分です。
実装
import sys
readline=sys.stdin.readline
import math
def Extended_Euclid(n,m):
stack=[]
while m:
stack.append((n,m))
n,m=m,n%m
if n>=0:
x,y=1,0
else:
x,y=-1,0
for i in range(len(stack)-1,-1,-1):
n,m=stack[i]
x,y=y,x-(n//m)*y
return x,y
N,K=map(int,readline().split())
P=[None]+list(map(int,readline().split()))
E=list(map(int,readline().split()))
mod=10**9+7
for i in range(1,N):
P[i]-=1
pow_N=pow(N,K,mod)
for i in range(N):
E[i]*=pow_N
E[i]%=mod
for i in range(N-1,0,-1):
E[P[i]]+=E[i]
E[P[i]]%=mod
if K==mod//2:
if all(E[i] in (0,1,mod-1) for i in range(N)):
child=[i for i in range(N) if P[i]==0 and E[i]]
if len(child)>=3:
ans=0
elif len(child)==2:
if E[0]:
ans=0
elif E[child[0]]!=E[child[1]]:
ans=0
else:
ans=1
elif len(child)==1:
if E[0]:
if E[0]==E[child[0]]:
ans=0
else:
ans=1
else:
if E[child[0]]==1:
ans=5
else:
ans=1
else:
if E[0]==1:
ans=1
elif E[0]==mod-1:
ans=5
else:
ans=0
else:
ans=-1
else:
K_inve=Extended_Euclid(K,mod-1)[0]%(mod-1)
for i in range(N):
E[i]=pow(E[i],K_inve,mod)
ans=E[0]
for i in range(1,N):
if P[i]==0:
ans-=E[i]
ans%=mod
print(ans)