Problem:


Score : 700 points

Problem Statement

There are N towns numbered 1, 2, \cdots, N.

Some roads are planned to be built so that each of them connects two distinct towns bidirectionally. Currently, there are no roads connecting towns.

In the planning of construction, each town chooses one town different from itself and requests the following: roads are built so that the chosen town is reachable from itself using one or more roads.

These requests from the towns are represented by an array P_1, P_2, \cdots, P_N. If P_i = -1, it means that Town i has not chosen the request; if 1 \leq P_i \leq N, it means that Town i has chosen Town P_i.

Let K be the number of towns i such that P_i = -1. There are (N-1)^K ways in which the towns can make the requests. For each way to make requests, find the minimum number of roads needed to meet all the requests, and print the sum of those (N-1)^K numbers, modulo (10^9+7).

Constraints

2 \leq N \leq 5000

P_i = -1 or 1 \leq P_i \leq N.

P_i \neq i

All values in input are integers.

Input

Input is given from Standard Input in the following format:

N
P_1 P_2 \cdots P_N

Output

For each way to make requests, find the minimum number of roads needed to meet all the requests, and print the sum of those (N-1)^K numbers, modulo (10^9+7).

Sample Input 1

4
2 1 -1 3

Sample Output 1

8

There are three ways to make requests, as follows:

Choose P_1 = 2, P_2 = 1, P_3 = 1, P_4 = 3. In this case, at least three roads - for example, (1,2),(1,3),(3,4) - are needed to meet the requests.

Choose P_1 = 2, P_2 = 1, P_3 = 2, P_4 = 3. In this case, at least three roads - for example, (1,2),(1,3),(3,4) - are needed to meet the requests.

Choose P_1 = 2, P_2 = 1, P_3 = 4, P_4 = 3. In this case, at least two roads - for example, (1,2),(3,4) - are needed to meet the requests.

Note that it is not mandatory to connect Town i and Town P_i directly.

The sum of the above numbers is 8.

Sample Input 2

2
2 1

Sample Output 2

1

There may be just one fixed way to make requests.

Sample Input 3

10
2 6 9 -1 6 9 -1 -1 -1 -1

Sample Output 3

527841



Sample input:
4
2 1 -1 3


Sample output:
8

Submission status: Wrong Answer

Solution code:
class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def members(self, x):
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

    def all_group_members(self):
        return {r: self.members(r) for r in self.roots()}

    def __str__(self):
        return '\n'.join('{}: {}'.format(r, self.members(r)) for r in self.roots())


n = int(input())

A = list(map(int, input().split()))
kosuu = A.count(-1)
print(kosuu)
uf = UnionFind(n)
for i, s in enumerate(A):
    if s != -1:
        uf.union(i, s-1)

print(uf)
B = uf.roots()
print(B)
num = 0
for s in B:
    num += uf.size(s)-1
mod = 10**9+7

ans = num*pow(n-1, kosuu, mod) % mod
for i, s in enumerate(A):
    if s == -1:
        ans = ans+(n-uf.size(i))*pow(n-1, kosuu-1, mod) % mod
print(ans)


# print(num)

Submission status: Wrong Answer

Solution code:
N=int(input())
P=list(map(int,input().split(" ")))
K=0
bl=[]
for i in range(len(P)):
    if P[i]==-1:
        K+=1
        bl+=[i]

R=len(P)-K
d2=0
for i in range(len(P)):
    if not(i in bl):
        if P[P[i]-1]==i+1:
            d2+=1

d2=int(d2/2)

bak=0
for i in range(len(P)):
    if not(i in bl):
        if P[P[i]-1]==-1:
            bak+=1

res=(((N-1)**(K-1))*((N-1)*K-bak)) % (10**9+7)

res+=(((N-1)**K)*(R-d2)) % (10**9+7)
print(res % (10**9+7))

Submission status: Wrong Answer

Solution code:
N = int(input())

Ps = list(map(int,input().split()))

K = Ps.count(-1)

ls = []

for i,P in enumerate(Ps):
  
  if P == -1:
    for n in range(len(Ps)):
      if n != i:
        ls.append((i,n))
  
  ls.append((i,P))
  
print(len(set(ls))%(10**9+7))

Submission status: Wrong Answer

Solution code:
input()
input()
print(1)

Submission status: Wrong Answer

Solution code:
N = int(input())
cityList = list(map(int, input().split()))
result = 0

for index, toCity in enumerate(cityList):

    if toCity == -1:
        continue

    if toCity > cityList[toCity - 1]:
        continue
        
    result += 1

print(result)

Submission status: Time Limit Exceeded

Solution code:
import sys
input = lambda : sys.stdin.readline().rstrip()
sys.setrecursionlimit(max(1000, 10**9))
write = lambda x: sys.stdout.write(x+"\n")


class UnionFindTree:
    def __init__(self, n):
        self.n = n
        self.parent = list(range(n))
        self.size = [1] * n

    def root(self, i):
        inter = set()
        while self.parent[i]!=i:
            inter.add(i)
            i = self.parent[i]
        r = i
        for i in inter:
            self.parent[i] = r
        return r

    def connect(self, i, j):
        if i==j:
            return
        ri = self.root(i)
        rj = self.root(j)
        if ri==rj:
            return
        if self.size[ri]<self.size[rj]:
            self.parent[ri] = rj
            self.size[rj] += self.size[ri]
        else:
            self.parent[rj] = ri
            self.size[ri] += self.size[rj]


n = int(input())
uf = UnionFindTree(n)
ps = list(map(int, input().split()))
k = 0
for i,p in enumerate(ps):
    if p<0:
        k += 1
        continue
    uf.connect(i,p-1)
for i in range(n):
    uf.root(i)
from collections import Counter, defaultdict
M = 10**9+7
ncount = Counter(uf.parent)
mcount = defaultdict(int)
for i,g in enumerate(uf.parent):
    if ps[i]==-1:
        mcount[g] += 1

ans = sum(v-1 for v in ncount.values())
ans *= pow(n-1, k, M)

for kk in ncount.keys():
    ncount[kk] -= mcount[kk]

for kk,mm in mcount.items():
    if mm==0:
        continue
    nn = ncount[kk]
    ans += (n - mm - nn)*pow(n-1, k-1, M)*mm
    ans %= M

nums = [nn+mcount[k] for k,nn in ncount.items() if mcount[k]]
l = len(nums)
dp = [[0] * l for _ in range(l)]
# dp[i][j] : i番目まで使うときの周期j+1の個数
if nums:
    dp[0][0] = nums[0]
    for i in range(1,l):
        dp[i][0] = dp[i-1][0] + nums[i]

    for i in range(1,l):
        for j in range(1,l):
            dp[i][j] += dp[i-1][j-1] * nums[i] * j + dp[i-1][j]

    for j in range(1, l):
        ans -= dp[-1][j]*pow(n-1, l-j-1, M)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
#!/usr/bin/env python3
import sys
sys.setrecursionlimit(10**8)
INF = float("inf")

MOD = 1000000007  # type: int


class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def members(self, x):
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

    def all_group_members(self):
        return {r: self.members(r) for r in self.roots()}

    def __str__(self):
        return '\n'.join('{}: {}'.format(r, self.members(r)) for r in self.roots())


def main():
    N = int(input())
    P = list(map(int, input().split()))

    # 確定している道でグルーピングする
    uf = UnionFind(N)
    K = 0
    undef = []                  # 要請未完ノードは記録
    for i, p in enumerate(P):
        if p != -1:
            uf.union(i, p-1)
        else:
            K += 1
            undef.append(i)

    # 階乗の事前計算
    factorial = [1]*K
    for i in range(K-1):
        factorial[i+1] = ((i+1)*factorial[i]) % MOD

    # 要請が未完のノードが属するグループサイズをメモ
    size_undef = []
    for u in undef:
        size_undef.append(uf.size(u))

    # dp[i, j]: i個の根付き木を見てj個を選んだ時の\sum{\prod_v S_v}
    # 「メモリ削減したdpとN^2確保するDP、どっちが早いんやろなあ」
    dp = [[0]*(K+1) for i in range(K+1)]
    dp[0][0] = 1
    for i in range(1, K+1):
        dp[i][0] = 1
        for j in range(i):
            dp[i][j+1] += dp[i-1][j]*size_undef[i-1] + dp[i-1][j+1]

    # ただしj=1個を選んだ時のものは特別にsize[i]-1を考える必要がある
    if K >= 1:
        tot = 0
        for s in size_undef:
            tot += s-1
        dp[K][1] = tot

    # cycleをカウントする
    cycle_count = 0
    for i in range(1, K+1):
        # i個のグループからなるサイクルがひとつできる
        buf = dp[K][i]*factorial[i-1]
        buf %= MOD
        buf *= pow(N-1, K-i, MOD)
        buf %= MOD
        cycle_count += buf

    # すでに固定サイクルがあれば数える
    cycle_count += (uf.group_count()-K)*pow(N-1, K, MOD)

    # 求めます
    ans = N*pow(N-1, K, MOD) - cycle_count
    ans %= MOD

    print(ans)

    return


if __name__ == '__main__':
    main()

Submission status: Time Limit Exceeded

Solution code:
#!/usr/bin/env python3
import sys
from functools import lru_cache
sys.setrecursionlimit(10**8)
INF = float("inf")

MOD = 1000000007  # type: int


@lru_cache(maxsize=None)
def factorial(n):
    if n == 0:
        return 1
    else:
        return (n*factorial(n-1)) % MOD


class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def members(self, x):
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

    def all_group_members(self):
        return {r: self.members(r) for r in self.roots()}

    def __str__(self):
        return '\n'.join('{}: {}'.format(r, self.members(r)) for r in self.roots())


def main():
    N = int(input())
    P = list(map(int, input().split()))

    # 確定している道でグルーピングする
    uf = UnionFind(N)

    K = 0
    undef = []
    for i, p in enumerate(P):
        if p != -1:
            uf.union(i, p-1)
        else:
            K += 1
            undef.append(i)

    groups = {}
    mapping = {}
    for i, r in enumerate(uf.roots()):
        groups[r] = [uf.size(r), True]
        mapping[i] = r
    for u in undef:
        groups[uf.find(u)][1] = False

    # 作り直し（むだっぽい）
    size_all = []
    looped = []
    size_undef = []
    for k, v in groups.items():
        size_all.append(v[0])
        looped.append(v[1])
        if v[1] == False:
            size_undef.append(v[0])
    # print(size_all)
    # print(looped)

    # dp[i, j]: i個の根付き木を見てj個を選んだ時の\sum{\prod_v S_v}
    dp = [0]*(K+1)
    dp[0] = 1
    dp_new = dp[:]
    for i in range(1, K+1):
        for j in range(i):
            dp_new[j+1] += dp[j]*size_undef[i-1]
        dp = dp_new[:]
        # print("dpしたときのdp", dp)

    # ただしj=1個を選んだ時のものは特別にsize[i]-1を考える必要がある
    if K >= 1:
        tot = 0
        for s in size_undef:
            tot += s-1
        dp[1] = tot
        # print("j==1のケースを修正後dp", dp)

    cycle_count = 0
    for i in range(1, K+1):
        # i個のグループからなるサイクルがひとつできる
        buf = dp[i]*factorial(i-1)
        buf %= MOD
        buf *= pow(N-1, K-i, MOD)
        buf %= MOD
        cycle_count += buf
        # print("数え上げ", dp[i], factorial(i-1), pow(N-1, K-i, MOD))
    # print(f"cycle_count, {cycle_count}")

    # すでに固定サイクルがあれば数える
    cycle_count += (len(groups)-K)*pow(N-1, K, MOD)

    # \sigma_{つなぎ方} N
    ans = N*pow(N-1, K, MOD)
    # print("sigma_{つなぎ方} N", ans)
    ans %= MOD
    ans = ans - cycle_count
    ans %= MOD

    print(ans)

    return


if __name__ == '__main__':
    main()

Submission status: Time Limit Exceeded

Solution code:
N=int(input())
A=list(map(int,input().split()))

MOD = 10**9 + 7
yet=[]
belongto=[-1]*N
group=[]
groupsize=[]

for i in range(N):
    if A[i]>0:
        if belongto[A[i]-1]<0 and belongto[i]<0:
            belongto[A[i]-1]=len(group)
            belongto[i]=len(group)
            group.append([i+1,A[i]])
        elif belongto[A[i]-1]>=0 and belongto[i]<0:
            belongto[i]=belongto[A[i]-1]
            group[belongto[A[i]-1]].append(i+1)
        elif belongto[A[i]-1]<0 and belongto[i]>=0:
            belongto[A[i]-1]=belongto[i]
            group[belongto[i]].append(A[i])
        elif belongto[A[i]-1]==belongto[i]:
            continue
        elif belongto[A[i]-1]>=0 and belongto[i]>=0:
            ori_group=belongto[A[i]-1]
            for j in group[belongto[A[i]-1]]:
                belongto[j-1]=belongto[i]
            group[belongto[i]]+=group[ori_group]
            group[ori_group]=[]
    elif A[i]<0:
        yet.append(i+1)
        if belongto[i]==-1:
            belongto[i]=len(group)
            group.append([i+1])
for i in range(len(group)):
    groupsize.append(len(group[i]))

    
if len(yet)>0:
    comb=[[0]*len(yet) for _ in range(len(yet))]
    comb[0][0]=groupsize[belongto[yet[0]-1]]-1

    for i in range(len(yet)-1):
        comb[i+1][0]=(comb[i][0]+groupsize[belongto[yet[i+1]-1]]-1)%MOD
        comb[i+1][1]=((comb[i][0]+i+1)*groupsize[belongto[yet[i+1]-1]])%MOD
        for j in range(1,i+1):
            comb[i+1][j]=(comb[i][j]+comb[i+1][j])%MOD
            comb[i+1][j+1]=(comb[i][j]*groupsize[belongto[yet[i+1]-1]]*(j+1))%MOD

groupnum =0
for agroup in group:
    if len(agroup)>0:
        groupnum+=1

total=(N-(groupnum-len(yet)))%MOD
if len(yet)==0:
    print(total)
else:
    for i in range(len(yet)):
        total = (total*(N-1))%MOD
        total=(total-comb[len(yet)-1][i])%MOD

    print(total)

Submission status: Time Limit Exceeded

Solution code:
from collections import deque

N = int(input())
P = list(map(int, input().split()))
MOD = 10 ** 9 + 7

a_b = {i:set() for i in range(N)}
b_a = {}
a_1 = deque()
for i, p in enumerate(P):
  if p == -1:
    a_1.append(i)
  else:
    j = p - 1
    if i in b_a:
      i = b_a[i]
    if j in b_a:
      j = b_a[j]
    if i != j:
      if i > j:
        i, j = j, i

      for k in a_b[j]:
        b_a[k] = i
      b_a[j] = i
      a_b[i] |= a_b[j]
      a_b[i].add(j)
      del a_b[j]
a_0 = {k for k, v in a_b.items() if len(v) == 0}
for x in a_0:
  del a_b[x]

ans = 0
que = deque()
que.append((a_b, b_a, a_0, list(a_1), 1))
while que:
  a_b, b_a, a_0, a_1, dup = que.pop()
  if a_1:
    i = a_1[0]
    a_1 = a_1[1:]
    if i in b_a:
      i = b_a[i]
    if i in a_0:
      a_0 = a_0 - {i}
    for j in a_b:
      dup2 = dup * (len(a_b[j]) + (0 if i == j or i in a_b[j] else 1)) % MOD

      if i != j:
        a_b2 = a_b.copy()
        b_a2 = b_a.copy()
        if i in a_b2:
          ii, jj = i, j
          if ii > jj:
            ii, jj = jj, ii

          for k in a_b2[jj]:
            b_a2[k] = ii
          b_a2[jj] = ii
          a_b2[ii] = a_b2[ii] | a_b2[jj]
          a_b2[ii].add(jj)
          del a_b2[jj]
        else:
          if i < j:
            for k in a_b2[j]:
              b_a2[k] = i
            b_a2[j] = i
            a_b2[i] = a_b2[j] | {j}
            del a_b2[j]
          else:
            b_a2[i] = j
            a_b2[j] = a_b2[j] | {i}
      else:
        a_b2, b_a2 = a_b, b_a

      que.append((a_b2, b_a2, a_0, a_1, dup2))
    if a_0:
      dup2 = dup * len(a_0) % MOD
      a_0 = a_0.copy()
      j = a_0.pop()
      a_b2 = a_b.copy()
      b_a2 = b_a.copy()
      if i in a_b2:
        if i < j:
          b_a2[j] = i
          a_b2[i] = a_b2[i] | {j}
        else:
          for k in a_b2[i]:
            b_a2[k] = j
          b_a2[i] = j
          a_b2[j] = a_b2[i] | {i}
          del a_b2[i]
      else:
        ii, jj = i, j
        if ii > jj:
          ii, jj = jj, ii
        b_a2[jj] = ii
        a_b2[ii] = {jj}

      que.append((a_b2, b_a2, a_0, a_1, dup2))
  else:
    ans = (ans + (sum(len(v) for v in a_b.values()) % MOD) * dup) % MOD

print(ans)

Submission status: Time Limit Exceeded

Solution code:
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 = [0] * (m + 1)  # 閉路のパターン数
dp[0] = 1
additional = 0  # 閉路を考慮せずに追加される辺数

for i in undefined:
    c = uft.get_size(i)
    for j in range(m - 1, -1, -1):
        dp[j + 1] = (dp[j + 1] + dp[j] * c) % MOD
    additional += n - c

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)

Submission status: Time Limit Exceeded

Solution code:
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 = [0] * (m + 1)  # 閉路のパターン数
dp[0] = 1
additional = 0  # 閉路を考慮せずに追加される辺数

for i in undefined:
    c = uft.get_size(i)
    for j in range(m - 1, -1, -1):
        dp[j + 1] += dp[j] * c
    additional += n - c

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)

Submission status: Time Limit Exceeded

Solution code:
from collections import deque

N = int(input())
P = list(map(int, input().split()))
MOD = 10 ** 9 + 7

a_b = {i:set() for i in range(N)}
b_a = {}
a_1 = deque()
for i, p in enumerate(P):
  if p == -1:
    a_1.append(i)
  else:
    j = p - 1
    if i in b_a:
      i = b_a[i]
    if j in b_a:
      j = b_a[j]
    if i != j:
      if i > j:
        i, j = j, i

      for k in a_b[j]:
        b_a[k] = i
      b_a[j] = i
      a_b[i] |= a_b[j]
      a_b[i].add(j)
      del a_b[j]

ans = 0
que = deque()
que.append((a_b, b_a, list(a_1), 1))
while que:
  a_b, b_a, a_1, dup = que.pop()
  if a_1:
    i = a_1[0]
    a_1 = a_1[1:]
    if i in b_a:
      i = b_a[i]
    for j in a_b:
      if i != j or a_b[j]:
        dup2 = dup * (len(a_b[j]) + (0 if i == j or i in a_b[j] else 1)) % MOD

        if i != j:
          a_b2 = a_b.copy()
          b_a2 = b_a.copy()
          ii, jj = i, j
          if ii > jj:
            ii, jj = jj, ii

          for k in a_b2[jj]:
            b_a2[k] = ii
          b_a2[jj] = ii
          a_b2[ii] = a_b2[ii] | a_b2[jj]
          a_b2[ii].add(jj)
          del a_b2[jj]
        else:
          a_b2, b_a2 = a_b, b_a

        que.append((a_b2, b_a2, a_1, dup2))
  else:
    ans = (ans + (sum(len(v) for v in a_b.values()) % MOD) * dup) % MOD

print(ans)

Submission status: Time Limit Exceeded

Solution code:
from collections import defaultdict, deque

N = int(input())
P = list(map(int, input().split()))

a_b = {i:set() for i in range(N)}
b_a = {}
a_1 = deque()
for i, p in enumerate(P):
  if p == -1:
    a_1.append(i)
  else:
    j = p - 1
    if i in b_a:
      i = b_a[i]
    if j in b_a:
      j = b_a[j]
    if i == j:
      continue
    if i > j:
      i, j = j, i

    for k in a_b[j]:
      b_a[k] = i
    b_a[j] = i
    a_b[i] |= a_b[j]
    a_b[i].add(j)
    del a_b[j]

ans = 0
que = deque()
que.append((a_b, b_a, a_1, 1))
while que:
  a_b, b_a, a_1, dup = que.pop()
  if a_1:
    a_1 = a_1.copy()
    i = a_1.popleft()
    if i in b_a:
      i = b_a[i]
    for j in a_b.keys():
      if i == j and len(a_b[j]) == 0:
        continue
      if (i == j) or (i in a_b[j]):
        dup2 = dup * len(a_b[j])
      else:
        dup2 = dup * (len(a_b[j]) + 1)

      ii, jj = i, j
      if ii != jj:
        a_b2 = {k:v.copy() for k, v in a_b.items()}
        b_a2 = b_a.copy()
        if ii > jj:
          ii, jj = jj, ii

        for k in a_b2[jj]:
          b_a2[k] = ii
        b_a2[jj] = ii
        a_b2[ii] |= a_b2[jj]
        a_b2[ii].add(jj)
        del a_b2[jj]
      else:
        a_b2, b_a2 = a_b, b_a

      que.append((a_b2, b_a2, a_1, dup2))
  else:
    ans += sum(len(a_b[x]) for x in a_b) * dup

print(ans)

Submission status: Time Limit Exceeded

Solution code:
import itertools
import copy

class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

N = int(input())
P = list(map(int, input().split()))

uf = UnionFind(N)
for i in range(N):
    if P[i] != -1:
        uf.union(i, P[i]-1)

result = 0
K = P.count(-1)
for comb in itertools.product(range(N), repeat=K):
    n = 0
    uf2 = copy.deepcopy(uf)
    for i in range(N):
        if P[i] == -1:
            if i == comb[n]:
                break
            uf2.union(i, comb[n])
            n += 1
    else:
        result += N-uf2.group_count()

print(result%(10**9+7))

Submission status: Time Limit Exceeded

Solution code:
from collections import defaultdict, deque

N = int(input())
P = list(map(int, input().split()))

a_b = {i:set() for i in range(N)}
b_a = {}
a_1 = deque()
for i, p in enumerate(P):
  if p == -1:
    a_1.append(i)
  else:
    j = p - 1
    if i in b_a:
      i = b_a[i]
    if j in b_a:
      j = b_a[j]
    if i == j:
      continue
    if i > j:
      i, j = j, i

    for k in a_b[j]:
      b_a[k] = i
    b_a[j] = i
    a_b[i] |= a_b[j]
    a_b[i].add(j)
    del a_b[j]

ans = 0
que = deque()
que.append((a_b, b_a, a_1, 1))
while que:
  a_b, b_a, a_1, dup = que.pop()
  if a_1:
    a_1 = a_1.copy()
    i = a_1.popleft()
    if i in b_a:
      i = b_a[i]
    for j in a_b.keys():
      if i == j and len(a_b[j]) == 0:
        continue
      a_b2 = {k:v.copy() for k, v in a_b.items()}
      b_a2 = b_a.copy()
      if (i == j) or (i in a_b[j]):
        dup2 = dup * len(a_b[j])
      else:
        dup2 = dup * (len(a_b[j]) + 1)

      ii, jj = i, j
      if ii != jj:
        if ii > jj:
          ii, jj = jj, ii

        for k in a_b2[jj]:
          b_a2[k] = ii
        b_a2[jj] = ii
        a_b2[ii] |= a_b2[jj]
        a_b2[ii].add(jj)
        del a_b2[jj]

      que.append((a_b2, b_a2, a_1, dup2))
  else:
    ans += sum(len(a_b[x]) for x in a_b) * dup

print(ans)

Submission status: Time Limit Exceeded

Solution code:
import itertools
import numpy as np
from scipy.sparse.csgraph import minimum_spanning_tree
from scipy.sparse import csr_matrix, coo_matrix, lil_matrix

def solve(P):
    #print(P)
    mat = np.zeros((len(P), len(P)))
    for i, j in enumerate(P):
        mat[i, j] = 1
        mat[j, i] = 1
    mat = mat[1:, 1:]
    mat = csr_matrix(mat)
    #print(mat)
    mst = minimum_spanning_tree(mat)
    return len(mst.data)


def main():
    N = int(input())
    P = [0]
    P.extend([int(c) for c in input().split(' ')])

    pos = [p[0] for p in enumerate(P) if p[1]==-1]

    ans = 0
    for t in itertools.product(range(1, N+1), repeat=len(pos)):
        assert len(t) == len(pos)
        Q = P.copy()
        flag = False
        for i, position in enumerate(pos):
            if t[i] == position:
                flag=True
                break
            else:
                Q[position] = t[i]
        if flag:
            continue
        ans += solve(Q)
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
from collections import defaultdict, deque, Counter
from heapq import heappush, heappop, heapify
import math
import bisect
import random
from itertools import permutations, accumulate, combinations, product
import sys
import string
from bisect import bisect_left, bisect_right
from math import factorial, ceil, floor, cos, radians, pi, sin
from operator import mul
from functools import reduce
from operator import mul
from functools import lru_cache

mod = 10 ** 9 + 7
sys.setrecursionlimit(2147483647)
INF = 10 ** 13
def LI(): return list(map(int, sys.stdin.buffer.readline().split()))
def I(): return int(sys.stdin.buffer.readline())
def LS(): return sys.stdin.buffer.readline().rstrip().decode('utf-8').split()
def S(): return sys.stdin.buffer.readline().rstrip().decode('utf-8')
def IR(n): return [I() for i in range(n)]
def LIR(n): return [LI() for i in range(n)]
def SR(n): return [S() for i in range(n)]
def LSR(n): return [LS() for i in range(n)]
def SRL(n): return [list(S()) for i in range(n)]
def MSRL(n): return [[int(j) for j in list(S())] for i in range(n)]


@lru_cache(maxsize=None)
def factorial(n):
    if n == 0:
        return 1
    else:
        return (n*factorial(n-1)) % mod


class UnionFind:
    def __init__(self, n):
        # 負  : 根であることを示す。絶対値はランクを示す
        # 非負: 根でないことを示す。値は親を示す
        self.table = [-1] * n
        self.size = [1] * n
        self.group_num = n

    def root(self, x):
        if self.table[x] < 0:
            return x
        else:
            self.table[x] = self.root(self.table[x])
            return self.table[x]

    def get_size(self, x):
        r = self.root(x)
        return self.size[r]

    def is_same(self, x, y):
        return self.root(x) == self.root(y)

    def union(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.size[r1] += self.size[r2]
            if d1 == d2:
                self.table[r1] -= 1
        else:
            self.table[r1] = r2
            self.size[r2] += self.size[r1]
        self.group_num -= 1

n = I()
P = LI()

fac = [1] * (n + 1)
inv = [1] * (n + 1)
for j in range(1, n + 1):
    fac[j] = fac[j-1] * j % mod


inv[n] = pow(fac[n], mod-2, mod)
for j in range(n-1, -1, -1):
    inv[j] = inv[j+1] * (j+1) % mod


def comb(n, r):
    if r > n or n < 0 or r < 0:
        return 0
    return fac[n] * inv[n - r] * inv[r] % mod


U = UnionFind(n)
for i in range(n):
    if P[i] != -1:
        U.union(P[i] - 1, i)

k = 0
L = []
for j in range(n):
    if P[j] == -1:
        L += [U.get_size(j)]
        k += 1

dp = [0]*(k+1)
dp[0] = 1
dp_new = dp[:]
for i in range(1, k+1):
    for j in range(i):
        dp_new[j+1] += dp[j]*L[i-1]
    dp = dp_new[:]


if k:
    dp[1] = sum(L) - len(L)

cycle_count = 0
for i in range(1, k + 1):
    # i個のグループからなるサイクルがひとつできる
    buf = dp[i] * factorial(i - 1)
    buf %= mod
    buf *= pow(n - 1, k - i, mod)
    buf %= mod
    cycle_count += buf


# すでに固定サイクルがあれば数える
cycle_count += (U.group_num - k) * pow(n - 1, k, mod)
ans = n * pow(n - 1, k, mod)
ans %= mod
ans = ans - cycle_count
ans %= mod
print(ans)