配列なんかでたまにある「左から可能な最大」「右から可能な最大」をそれぞれ求めて、minを取ったのが全体で可能な最大、というのを木でやる感じ。
葉でない頂点は、1個または2個の子を持つ。
○ ○ 深さ d-1 ／ ／＼ ○ ○ ○ d
なので、ある深さ $d$ で作れる頂点数は、$d-1$ の葉でない頂点数の1倍～2倍となる。
適当に子が1つの頂点と2つの頂点の個数を調整すれば、この範囲内なら深さ $d$ の頂点数を自由に決められる。
なので、深さ $d$ の頂点数をこの範囲内にどうやったってできないというのが、不可能な場合である。
「少なすぎてできない」というのはあり得なくて、やろうと思えばずっと子が1つだけの頂点を繋げていける。
子の数
○ 0
／
○ 0
／
○ 0
／
...
「多すぎてできない」のを考える。つまり、深さ $d$ で作れる頂点数が $V_d$ なのに、それを超える $A_d$ 個の葉を作れと言われても、無理である。
根から、できる最大の頂点数を調べていく。
Ai 頂点数の最大 葉でない頂点の最大
0 ○ 1
／＼ ↙ x2
0 ○ ○ 2 → -Ai 2
/\ /\
1 ○○○● ←1個は葉 4 3
0 OOOOOO 6 6
2 OO....OO 12
こんな感じで、2倍して、葉なる頂点数 $A_i$ を引いて、とすればよい。
この過程で、頂点数の最大が $A_i$ より少なくなったらアウト。アウトにならなければ、可能。
ただ、これはあくまで根から考えた最大であって、このままでは $A_N$ との整合が取れない。
よって、次に下から考える。
さっきの計算を逆に考えると、ある深さ $d$ の頂点数がわかっているとすると、$d-1$ の頂点数はこれを1～2で割った数となる。
なるべく頂点数を多く保ちたいのであれば、1で割る（つまり、全ての親を別々の頂点にする）のがよい。
○ ○ ○ ○ ○ 深さ d-1 頂点数 5 ／ ／ ／ ／ ／ ○ ○ ○ ○ ○ 深さ d 頂点数 5
これをもとに、下から可能な頂点数の最大を調べる。その際、根から求めた頂点数を超えないようにする。
※木を上下ひっくり返す
根から考えた 下から考えた 上下統合した
Ai 頂点数の最大 葉でない頂点数の最大 頂点数の最大
2 ● ● 12 2
\ \ ↙------'
0 ○ ○ 6 2 +Ai → 2
\ \
1 ○ ○ ● 4 2 3
\ / /
0 ○ ○ 2 3 3 → 2
＼／ `→ 根からの最大を超えてしまう場合は、minをとる
0 ○ 1 2 2 → 1
-----
ここの合計が、答えとなる→ 10
def solve(n, aaa):
if n == 0:
if aaa[0] == 1:
return 1
else:
return -1
if aaa[0] != 0:
return -1
hi = 1
hi_enable = [1]
for a in aaa[1:]:
if hi * 2 < a:
return -1
hi = hi * 2 - a
hi_enable.append(hi)
vn = aaa[n]
ans = vn
for i in range(n - 1, -1, -1):
hi = hi_enable[i]
a = aaa[i]
vn = min(vn, hi) + a
ans += vn
return ans
n = int(input())
aaa = list(map(int, input().split()))
print(solve(n, aaa))
'-1' なら未定、それ以外なら指定した都市を示す$i$ と $P_i$ が、他の都市の指定により既に道路を通って行き来できるなら新たな道路を建設する必要は無い。
なので、既に指定済みの都市に関しては、Union-Find木で連結成分を管理する。（連結の頂点数も求められるようにしておく）
各都市からは最大1つしか辺を伸ばせないため、ある連結成分に $P_i=-1$（未指定）の都市が含まれているとしても、多くとも1つである。 （連結な $N$ 頂点のグラフの最小辺数は $N-1$。“-1” からはこの時点ではまだ辺を伸ばさないため、2個以上含まれると辺数が $N-2$ 以下になる）
指定済みの都市を繋げた結果、こんな感じになる。この時点の6本の道路は、残りをどのように決めるとしても必要になる。
○-○-○ 頂点数 3 ○: 指定済み ●-○-○-○ 頂点数 4 ●: 未指定 ●-○ 頂点数 2 ● 頂点数 1
ここから、●を決めた時の新設道路数を考えていく。
基本的には、「自身が含まれる連結成分内の頂点を指定すると0」「含まれない頂点を指定すると+1」となる。
しかし、ここでも他の●の決め方によっては、指定した頂点が既に連結かも知れない。
❶-○-○-○ ❶が、❷の連結成分の頂点を指定した場合 ↓ ❷-○ →❷を決めるとき、自身に加え、❶の頂点も既に連結成分
2個だけならまだしも、3個、4個をまたいで広く連結になっているかも知れない。
❶-○-○-○ ↓ ↑ ❸ ❷-○ ❸は、既に❶とも❷とも連結
これを、上手く除いて考えなければならない。この方法がわからなかった。
この問題のように、各頂点から「1本のみ」辺を伸ばすとき、無駄な辺ができるとしても、各連結成分内に多くとも1本である。 （無駄な辺ができる場合、閉路ができる）
●←● ●⇄●←● ↓↗ ●
この「閉路」を固定して考えると、閉路ごとに独立に考えられて上手くいく。
元の問題に戻って、❶❷が互いに互いの連結成分を指定し合う閉路を考える。 以下、|❶| で❶の連結成分数を示すとする。
❶-○-○-○ ↓ ↑ ○-❷ ❸
この場合の選び方は、$|❶| \times |❷|＝8$ 通り。（あとは❸の選び方が $N-1$ 通り）
同様に❶❸、❷❸の閉路も考えると、経由する連結成分が2個である閉路は以下になる。
また、長さ3の、❶❷❸全てを巡る閉路は、巡る順番もパターンに含まれる。
❶-○-○-○ ○-○-○-❶ ↓ ↑ ↑ ↓ ○-❷ → ❸ ❷-○ ← ❸
よって、閉路の部分だけを考えるとパターン数は $|❶||❷||❸| \times 2!$ 通りある。
このパターン数は、経由成分数ごとに、以下のようなDPでまとめて計算できる。
経由成分数 0 1 2 3
1
↘ |❶|倍して加算
❶ 1 |❶|
↘ ↘
❷ 1 |❶|+|❷| |❶||❷|
↘ ↘ ↘
❸ 1 |❶|+|❷|+|❸| |❶||❷| |❶||❷||❸|
+|❶||❸|
+|❷||❸|
経由成分数 $l$ の閉路は $DP[l] \times (l-1)!$ 通りあり、同じ閉路が他の未指定頂点の決め方 $(N-1)^{k-l}$ 通りで1個ずつ生じる。
他の未指定頂点の中で閉路ができたとしても、「今、着目中の閉路が生じるパターン数」には関係ない。別々に数えることができる。
よって、各未指定頂点につき、とりあえず閉路は考えないで +1 して数えた後、上記の閉路を引けばよい。
ただし、経由成分数が1つ（自身の連結成分に繋ぐ）場合、この問題では自分自身は選べないため、実際に選べる頂点は1ずつ少ない点に注意。
まとめると、
import numpy as np
class UnionFind:
def __init__(self, n):
self.table = [-1] * n
def _root(self, x):
stack = []
tbl = self.table
while tbl[x] >= 0:
stack.append(x)
x = tbl[x]
for y in stack:
tbl[y] = x
return x
def find(self, x, y):
return self._root(x) == self._root(y)
def unite(self, x, y):
r1 = self._root(x)
r2 = self._root(y)
if r1 == r2:
return
d1 = self.table[r1]
d2 = self.table[r2]
if d1 <= d2:
self.table[r2] = r1
self.table[r1] += d2
else:
self.table[r1] = r2
self.table[r2] += d1
def get_size(self, x):
return -self.table[self._root(x)]
n = int(input())
ppp = list(map(int, input().split()))
MOD = 10 ** 9 + 7
uft = UnionFind(n)
base = 0
undefined = []
for i, p in enumerate(ppp):
if p == -1:
undefined.append(i)
continue
p -= 1
if not uft.find(i, p):
base += 1
uft.unite(i, p)
if len(undefined) == 0:
print(base % MOD)
exit()
elif len(undefined) == 1:
c = uft.get_size(undefined[0])
others = n - c
print((base * (n - 1) + others) % MOD)
exit()
m = len(undefined)
dp = np.zeros(m + 1, dtype=np.int64) # 閉路のパターン数
dp[0] = 1
additional = 0 # 閉路を考慮せずに追加される辺数
for i in undefined:
c = uft.get_size(i)
dp[1:] += dp[:-1] * c
dp %= MOD
additional += n - c
dp = dp.tolist()
duplicated = 0
pat = pow(n - 1, m, MOD)
inv = pow(n - 1, MOD - 2, MOD)
loop_permutation = 1
loop_other_pattern = pat * inv * inv % MOD
for loop_size in range(2, m + 1):
duplicated = (duplicated + dp[loop_size] * loop_permutation * loop_other_pattern) % MOD
loop_other_pattern = loop_other_pattern * inv % MOD
loop_permutation = loop_permutation * loop_size % MOD
ans = (base * pat + additional * pat * inv - duplicated) % MOD
print(ans)