Submission #13919692


Source Code Expand

Copy
import sys
import numpy as np

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

MOD = 10**9 + 7

N, *P = map(int, read().split())

class UnionFind:
    def __init__(self, N):
        self.root = list(range(N))
        self.size = [1] * (N)
        self.has_loop = [False] * N

    def __getitem__(self, x):
        root = self.root
        while root[x] != x:
            root[x] = root[root[x]]
            x = root[x]
        return x

    def merge(self, x, y):
        x = self[x]
        y = self[y]
        if x == y:
            self.has_loop[x] = True
            return False
        sx, sy = self.size[x], self.size[y]
        if sx < sy:
            x, y = y, x
            sx, sy = sy, sx
        self.root[y] = x
        self.size[x] += sy
        self.has_loop[x] |= self.has_loop[y]
        return True

uf = UnionFind(N + 1)
for i, x in enumerate(P, 1):
    if x == -1:
        continue
    uf.merge(i, x)

loop = 0
sizes = []
for v in range(1, N + 1):
    if v != uf[v]:
        continue
    if uf.has_loop[v]:
        loop += 1
    else:
        sizes.append(uf.size[v])

S = len(sizes)

f = np.zeros(S + 10, np.int64)
f[0] = 1
for x in sizes:
    f[1:] += x * f[:-1]
    f %= MOD
f[1] -= S

fact = np.empty(N + 1, np.int64)
fact[0] = 1
for n in range(1, N + 1):
    fact[n] = n * fact[n - 1] % MOD
power = np.empty(N + 1, np.int64)
power[0] = 1
for n in range(1, N + 1):
    power[n] = (N - 1) * power[n - 1] % MOD

coef = fact[:S] * power[:S][::-1] % MOD

x = (f[1:S+1] * coef % MOD).sum() % MOD
x += loop * pow(N-1, S, MOD) % MOD

x = N * pow(N-1, S, MOD) - x
x %= MOD
print(x)

Submission Info

Submission Time
Task D - Urban Planning
User maspy
Language Python (3.8.2)
Score 700
Code Size 1735 Byte
Status AC
Exec Time 389 ms
Memory 27972 KB

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 700 / 700
Status
AC × 3
AC × 37
Set Name Test Cases
Sample s1.txt, s2.txt, s3.txt
All 01.txt, 02.txt, 03.txt, 04.txt, 05.txt, 06.txt, 07.txt, 08.txt, 09.txt, 10.txt, 11.txt, 12.txt, 13.txt, 14.txt, 15.txt, 16.txt, 17.txt, 18.txt, 19.txt, 20.txt, 21.txt, 22.txt, 23.txt, 24.txt, 25.txt, 26.txt, 27.txt, 28.txt, 29.txt, 30.txt, 31.txt, 32.txt, 33.txt, 34.txt, s1.txt, s2.txt, s3.txt
Case Name Status Exec Time Memory
01.txt AC 107 ms 26936 KB
02.txt AC 104 ms 27004 KB
03.txt AC 105 ms 27060 KB
04.txt AC 103 ms 26792 KB
05.txt AC 107 ms 27228 KB
06.txt AC 106 ms 26792 KB
07.txt AC 107 ms 26996 KB
08.txt AC 105 ms 26980 KB
09.txt AC 105 ms 27016 KB
10.txt AC 104 ms 26752 KB
11.txt AC 107 ms 27248 KB
12.txt AC 111 ms 27296 KB
13.txt AC 106 ms 27252 KB
14.txt AC 117 ms 27652 KB
15.txt AC 118 ms 27688 KB
16.txt AC 384 ms 27520 KB
17.txt AC 389 ms 27412 KB
18.txt AC 387 ms 27448 KB
19.txt AC 190 ms 27260 KB
20.txt AC 231 ms 27640 KB
21.txt AC 270 ms 27560 KB
22.txt AC 285 ms 27272 KB
23.txt AC 304 ms 27412 KB
24.txt AC 317 ms 27388 KB
25.txt AC 326 ms 27272 KB
26.txt AC 199 ms 27232 KB
27.txt AC 161 ms 27972 KB
28.txt AC 255 ms 27660 KB
29.txt AC 150 ms 27680 KB
30.txt AC 154 ms 27736 KB
31.txt AC 151 ms 27496 KB
32.txt AC 248 ms 27804 KB
33.txt AC 243 ms 27656 KB
34.txt AC 196 ms 27396 KB
s1.txt AC 102 ms 27028 KB
s2.txt AC 106 ms 27212 KB
s3.txt AC 105 ms 27008 KB