Harryxun/testdata huggingface dataset:
Buggy python submission (verdict: Runtime Error):
from fractions import gcd

N=int(input())
d={}
#l=[]
MOD = 1000000007
nz=0
for _ in range(N):
    a,b=map(int,input().split())
    if a==0 and b==0:
        nz+=1
        #p=(0,0)
    elif a==0:
        p=(0,1)
    elif b==0:
        p=(1,0)
    else:
        c=gcd(a,b)
        if b > 0:
            p=(a//c,b//c)
        else: # b < 0
            p=(-a//c,-b//c)
    
    #l.append(p)
    if p in d:
        d[p]+=1
    else:
        d[p] = 1

ans=1
#for p in l:
while len(d) > 0:
    p,np=d.popitem()
    if p==(1,0):
        q=(0,1)
    elif p==(0,1):
        q=(1,0)
    else:
        pa,pb=p
        if pa > 0:
            q=(-pb,pa)
        else: #pa < 0
            q=(pb,-pa)
    if q in d:
        nq=d.pop(q)
    else:
        nq=0
    
    ans *= ((2**np + 2**nq - 1) % MOD)
    ans %= MOD

ans -= 1
ans += nz
ans %= MOD
print(ans)

Fixed submission:
from fractions import gcd

N=int(input())
d={}
#l=[]
MOD = 1000000007
nz=0
for _ in range(N):
    a,b=map(int,input().split())
    if a==0 and b==0:
        nz+=1
        #p=(0,0)
        continue
    elif a==0:
        p=(0,1)
    elif b==0:
        p=(1,0)
    else:
        c=gcd(a,b)
        a,b=a//c,b//c
        if b > 0:
            p=(a,b)
        else: # b < 0
            p=(-a,-b)
    
    #l.append(p)
    if p in d:
        d[p]+=1
    else:
        d[p] = 1

ans=1
#for p in l:
while len(d) > 0:
    p,np=d.popitem()
    if p==(1,0):
        q=(0,1)
    elif p==(0,1):
        q=(1,0)
    else:
        pa,pb=p
        if pa > 0:
            q=(-pb,pa)
        else: #pa < 0
            q=(pb,-pa)
    if q in d:
        nq=d.pop(q)
    else:
        nq=0
    
    ans *= ((2**np + 2**nq - 1) % MOD)
    ans %= MOD

ans -= 1
ans += nz
ans %= MOD
print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
N=int(input())
As=list(map(int,input().split()))

cnt=0
for i in range(N):
    if i + As[i]>N:
        continue
    for j in range(i+As[i],N):
        if j-i == As[i] + As[j]:
            cnt+=1
print(cnt)

Fixed submission:
N=int(input())
As=list(map(int,input().split()))

d1={}
for i,e in enumerate(As):
    d=i-e
    if d > 0:
        if d in d1.keys():
            d1[d]+=1
        else:
            d1[d]=1
cnt=0
for i,e in enumerate(As):
    if (i+e) in d1.keys():
        cnt += d1[i+e]
print(cnt)

Buggy python submission (verdict: Time Limit Exceeded):
N=int(input())
As=list(map(int,input().split()))

cnt=0
for i in range(N):
    for j in range(i,N):
        if j-i == As[i] + As[j]:
            cnt+=1
print(cnt)

Fixed submission:
N=int(input())
As=list(map(int,input().split()))

d1={}
for i,e in enumerate(As):
    d=i-e
    if d > 0:
        if d in d1.keys():
            d1[d]+=1
        else:
            d1[d]=1
cnt=0
for i,e in enumerate(As):
    if (i+e) in d1.keys():
        cnt += d1[i+e]
print(cnt)
Buggy python submission (verdict: Time Limit Exceeded):
A,B,N = map(int, input().split())
ans = 0
for i in range(10**7):
    if i > N-1:
        break
    
    tmp = (A*(i+1))//B-A*((i+1)//B)
    if ans < tmp:
        ans = tmp
print(ans)

Fixed submission:
A,B,N = map(int, input().split())

x = min(B-1,N)

print((A*x)//B -A*(x//B))
Buggy python submission (verdict: Runtime Error):
import collections
import bisect

def f(v):
    #state[v] = 1
    a=As[v]
    idx = bisect.bisect_left(D,a)
    stack.append( (idx,D[idx]) ) # push
    D[idx] = a
    
    ans[v] = bisect.bisect_left(D,float('inf'))

    for u in V[v]:
        # if state[u] == 1: 木構造なので不要
        len_pre=len(stack)
        
        f(u)

        # 巻き戻し
        while len(stack) > len_pre:
            i_pre,d_pre=stack.pop()
            D[i_pre] = d_pre



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

D=[float('inf')]*N
stack=collections.deque([])
ans=[0]*N
#state=[0]*N

V=[[] for _ in range(N)]
for _ in range(N-1):
    u,v=map(int,input().split())
    V[u-1].append(v-1)

f(0)

print(*ans,sep='\n')

Fixed submission:
import collections
import bisect
import sys

sys.setrecursionlimit(10 ** 8)

def f(v):
    state[v] = 1
    a=As[v]
    idx = bisect.bisect_left(D,a)
    stack.append( (idx,D[idx]) ) # push
    D[idx] = a
    
    ans[v] = bisect.bisect_left(D,float('inf'))

    for u in V[v]:
        if state[u] == 0:
            f(u)

    # 巻き戻し
    i_pre,d_pre=stack.pop()
    D[i_pre] = d_pre



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

D=[float('inf')]*N
stack=collections.deque([])
ans=[0]*N
state=[0]*N

V=[[] for _ in range(N)]
for _ in range(N-1):
    u,v=map(int,input().split())
    V[u-1].append(v-1)
    V[v-1].append(u-1)

f(0)

print(*ans,sep='\n')
Buggy python submission (verdict: Time Limit Exceeded):
S=input()

N=len(S)
d={}
for i in range(2019):
    d[i]=0

d[0] = 1

for i in range(N):
    s=S[N-1-i:]
    d[int(s)%2019]+=1


ans = 0
for k in d.keys():
    v = d[k]
    if v > 1:
        ans += v*(v-1)//2
print(ans)

Fixed submission:
S=input()

N=len(S)
d={}
for i in range(2019):
    d[i]=0

d[0] = 1

mod=0
dig=1

for i in range(N):
    #s=S[N-1-i:]
    #d[int(s)%2019]+=1
    mod += dig * int(S[N-1-i])
    mod %= 2019
    d[mod] += 1
    dig *= 10
    dig %= 2019


ans = 0
for k in d.keys():
    v = d[k]
    if v > 1:
        ans += v*(v-1)//2
print(ans)

Buggy python submission (verdict: Time Limit Exceeded):
S=input()
N=len(S)

l=[]

if N < 4:
    ans=0
else:
    ans = 0
    for i in range(N-3):
        for j in range(i+3,N):
            s=S[i:j+1]
            if s in l:
                ans +=1
            elif (int(s))%2019==0:
                ans +=1
                l.append(s)
print(ans)

Fixed submission:
S=input()

N=len(S)
d={}
for i in range(2019):
    d[i]=0

d[0] = 1

mod=0
dig=1

for i in range(N):
    #s=S[N-1-i:]
    #d[int(s)%2019]+=1
    mod += dig * int(S[N-1-i])
    mod %= 2019
    d[mod] += 1
    dig *= 10
    dig %= 2019


ans = 0
for k in d.keys():
    v = d[k]
    if v > 1:
        ans += v*(v-1)//2
print(ans)

Buggy python submission (verdict: Time Limit Exceeded):
S=input()
N=len(S)

if N < 4:
    ans=0
else:
    ans = 0
    for i in range(N-3):
        for j in range(i+3,N):
            n=int(S[i:j+1])
            if n%2019==0:
                ans +=1
print(ans)

Fixed submission:
S=input()

N=len(S)
d={}
for i in range(2019):
    d[i]=0

d[0] = 1

mod=0
dig=1

for i in range(N):
    #s=S[N-1-i:]
    #d[int(s)%2019]+=1
    mod += dig * int(S[N-1-i])
    mod %= 2019
    d[mod] += 1
    dig *= 10
    dig %= 2019


ans = 0
for k in d.keys():
    v = d[k]
    if v > 1:
        ans += v*(v-1)//2
print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
from collections import Counter


N=int(input())
l=['']*N
for i in range(N):
    l[i]=input()

c = Counter(l).most_common()
top=c.pop(0)
ans=[top[0]]
freq=top[1]

while len(c) > 0:
    top=c.pop(0)
    if top[1] == freq:
        ans.append(top[0])
    else:
        break

ans.sort()

for e in ans:
    print(e)

Fixed submission:
from collections import Counter


N=int(input())
l=['']*N
for i in range(N):
    l[i]=input()

l.sort()

c=1
c_crt=1
pre=""
ans=[]
for e in l+[""]:
    if e == pre:
        c_crt+=1
    else:
        if c_crt > c:
            ans=[pre]
            c=c_crt
        elif c_crt == c:
            ans.append(pre)
        c_crt=1
    pre=e
for e in ans:
    print(e)

Buggy python submission (verdict: Time Limit Exceeded):
from collections import Counter

l=[]
N=int(input())
for _ in range(N):
    l.append(input())

c = Counter(l).most_common()
top=c.pop(0)
ans=[top[0]]
freq=top[1]

while len(c) > 0:
    top=c.pop(0)
    if top[1] == freq:
        ans.append(top[0])
    else:
        break

ans.sort()

for e in ans:
    print(e)

Fixed submission:
from collections import Counter


N=int(input())
l=['']*N
for i in range(N):
    l[i]=input()

l.sort()

c=1
c_crt=1
pre=""
ans=[]
for e in l+[""]:
    if e == pre:
        c_crt+=1
    else:
        if c_crt > c:
            ans=[pre]
            c=c_crt
        elif c_crt == c:
            ans.append(pre)
        c_crt=1
    pre=e
for e in ans:
    print(e)
Buggy python submission (verdict: Time Limit Exceeded):
N,K,Q=[int(e) for e in input().split(' ')]
lA=[int(input()) for _ in range(Q)]

for i in range(N):
    if Q-lA.count(i+1) >= K:
        print('No')
    else:
        print('Yes')

Fixed submission:
N,K,Q=[int(e) for e in input().split(' ')]

c=[0]*N
for _ in range(Q):
    c[int(input())-1]+=1

for i in range(N):
    if Q-c[i] >= K:
        print('No')
    else:
        print('Yes')
Buggy python submission (verdict: Runtime Error):
A,B=[int(e) for r in input().split(' ')]
print(max([A+B,A-B,A*B]))

Fixed submission:
A,B=[int(e) for e in input().split(' ')]
print(max([A+B,A-B,A*B]))
Buggy python submission (verdict: Runtime Error):
A,B=[int(e) for e in input()]
print(max([A+A-1,A+B,B+B-1]))



Fixed submission:
A,B=[int(e) for e in input().split(' ')]
print(max([A+A-1,A+B,B+B-1]))
Buggy python submission (verdict: Runtime Error):
N=int(input())
dataH=[int(e) for e in input().split(' ')]
Hmax=max(dataH)

dataH.reverse()
print(len(dataH)-dataH.index(max))


Fixed submission:
N=int(input())
dataH=[int(e) for e in input().split(' ')]
c=0
for i,H in enumerate(dataH):
  c=c+1 if max(dataH[:i+1])==H else c
print(c)
Buggy python submission (verdict: Time Limit Exceeded):
import re
N,Q=[int(e) for e in input().split(' ')]
S=input()
qs=[]
for i in range(Q):
    qs.append([int(e) for e in input().split(' ')])

sts=[]
for i in range(N):
    if i>0:
        if S[i-1:i+1]=='AC':
            sts.append(i)

for q in qs:
    print(len([1 for e in sts if e>=q[0] and e<q[1]]))
#    print(len(re.findall('AC',substr)))

Fixed submission:
import re
N,Q=[int(e) for e in input().split(' ')]
S=input()
qs=[]
for i in range(Q):
    qs.append([int(e) for e in input().split(' ')])

sts=[0]*N
for i in range(N):
    if i>0:
        if S[i-1:i+1]=='AC':
            sts[i]=sts[i-1]+1
        else:
            sts[i]=sts[i-1]
for q in qs:
    print(sts[q[1]-1]-sts[q[0]-1])

Buggy python submission (verdict: Time Limit Exceeded):
import re
N,Q=[int(e) for e in input().split(' ')]
S=input()

qs=[]
for i in range(Q):
    qs.append([int(e) for e in input().split(' ')])

for q in qs:
    substr=S[q[0]-1:q[1]]
    print(len(re.findall('AC',substr)))

Fixed submission:
import re
N,Q=[int(e) for e in input().split(' ')]
S=input()
qs=[]
for i in range(Q):
    qs.append([int(e) for e in input().split(' ')])

sts=[0]*N
for i in range(N):
    if i>0:
        if S[i-1:i+1]=='AC':
            sts[i]=sts[i-1]+1
        else:
            sts[i]=sts[i-1]
for q in qs:
    print(sts[q[1]-1]-sts[q[0]-1])
Buggy python submission (verdict: Runtime Error):
istr=input()
for ch in istr:
  if ch == '0':
    c0+=1
  else:
    c1+=1
print(min(c0,c1)*2)
      

Fixed submission:
istr=input()
c0=0
c1=0
for ch in istr:
    if ch == '0':
        c0+=1
    else:
        c1+=1
print(min(c0,c1)*2)

Buggy python submission (verdict: Runtime Error):
def crmv(c,istr):
    if len(istr)<=1:
        return c
    if not('1' in istr):
        return c

    rv=0
    for i in range(len(istr)-1):
        if istr[i] != istr[i+1]:
            tstr=istr[:i]+istr[i+1+1:]
            r=crmv(c+2,tstr)
            if r>rv:
                rv=r
    return rv

istr=input()
print(crmv(0,istr))

Fixed submission:
istr=input()
c0=0
c1=0
for ch in istr:
    if ch == '0':
        c0+=1
    else:
        c1+=1
print(min(c0,c1)*2)

Buggy python submission (verdict: Runtime Error):
cmax=0
def crmv(c,istr):
    if len(istr)<=1:
        return c
    if not('1' in istr) or not('0' in istr):
        return c

    rv=0
    for i in range(len(istr)-1):
        if istr[i] != istr[i+1]:
            tstr=istr[:i]+istr[i+1+1:]
            r=crmv(c+2,tstr)
            if r>rv:
                rv=r
    return rv

istr=input()
print(crmv(0,istr))

Fixed submission:
istr=input()
c0=0
c1=0
for ch in istr:
    if ch == '0':
        c0+=1
    else:
        c1+=1
print(min(c0,c1)*2)

Buggy python submission (verdict: Runtime Error):
def crmv(c,istr):
    if len(istr)<=1:
        return c
    '''
    lidx=[]
    for i in range(len(istr)-1):
        if istr[i] != istr[i+1]:
            lidx.append(i)
    if len(lidx)==0:
        return c

    rv=0
    for i in lidx:
        tstr=istr[:i]+istr[i+1+1:]
        r=crmv(c+2,tstr)
        if r>rv:
            rv=r
    '''
    rv=-1
    for i in range(len(istr)-1):
        if istr[i] != istr[i+1]:
            tstr=istr[:i]+istr[i+1+1:]
            r=crmv(c+2,tstr)
            if r>rv:
                rv=r
    if rv == -1:
        return c
            

    return rv

istr=input()
print(crmv(0,istr))

Fixed submission:
istr=input()
c0=0
c1=0
for ch in istr:
    if ch == '0':
        c0+=1
    else:
        c1+=1
print(min(c0,c1)*2)

Buggy python submission (verdict: Time Limit Exceeded):
def crmv(c,istr):
    if len(istr)<=1:
        return c
    
    lidx=[]
    for i in range(len(istr)-1):
        if istr[i] != istr[i+1]:
            lidx.append(i)
    if len(lidx)==0:
        return c

    rv=0
    for i in lidx:
        tstr=istr[:i]+istr[i+1+1:]
        r=crmv(c+1,tstr)
        if r>rv:
            rv=r
    return rv

istr=input()

print(crmv(0,istr)*2)

Fixed submission:
istr=input()
c0=0
c1=0
for ch in istr:
    if ch == '0':
        c0+=1
    else:
        c1+=1
print(min(c0,c1)*2)
Buggy python submission (verdict: Time Limit Exceeded):
# cook your dish here
n=int(input())
a=[int(x) for x in input().split()]
ans=1  
sorted(a)
for i in range(n):
    ans*=a[i]
    if ans==0:
        break
    if ans>(int)(1e18):
        ans=-1
print(ans)

Fixed submission:
# cook your dish here
n=int(input())
a=[int(x) for x in input().split()]
ans=1  
mod=int(1e18)
for i in range(n):
    if a[i]==0:
        ans=0 
        break 
for i in range(n):
    ans*=a[i]
    if ans>mod:
        ans=-1 
        break 
    if ans==0:
        break
print(ans)

Buggy python submission (verdict: Time Limit Exceeded):
# cook your dish here
n=int(input())
a=[int(x) for x in input().split()]
ans=1 
for i in range(n):
    ans*=a[i] 
if ans>(int)(1e18):
    ans=-1
print(ans)

Fixed submission:
# cook your dish here
n=int(input())
a=[int(x) for x in input().split()]
ans=1  
mod=int(1e18)
for i in range(n):
    if a[i]==0:
        ans=0 
        break 
for i in range(n):
    ans*=a[i]
    if ans>mod:
        ans=-1 
        break 
    if ans==0:
        break
print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
import sys
import io, os
input = sys.stdin.buffer.readline
#input = io.BytesIO(os.read(0,os.fstat(0).st_size)).readline

import math
N = 10**6+50

def main():

    n = int(input())
    A = list(map(int, input().split()))

    spf = [-1]*(N+1)
    for i in range(N+1):
        spf[i] = i
    for i in range(2, int(math.sqrt(N))+1):
        if spf[i] == i:
            for j in range(i*2, N+1, i):
                if spf[j] == j:
                    spf[j] = i

    S = set()
    flag = False
    for i, a in enumerate(A):
        s = set()
        while a != 1:
            p = spf[a]
            if p in S:
                #print('setwise coprime')
                #exit()
                flag = True
            s.add(p)
            a //= p
        S = S | s
        if i == 0:
            T = s
        else:
            T = T & s

    #print(S)
    #print(T)
    if len(T) != 0:
        print('not coprime')
    elif flag:
        print('setwise coprime')
    else:
        print('pairwise coprime')

if __name__ == '__main__':
    main()


Fixed submission:
import sys
import io, os
input = sys.stdin.buffer.readline
#input = io.BytesIO(os.read(0,os.fstat(0).st_size)).readline

import math
N = 10**6+50

def main():

    n = int(input())
    A = list(map(int, input().split()))

    spf = [-1]*(N+1)
    for i in range(N+1):
        spf[i] = i
    for i in range(2, int(math.sqrt(N))+1):
        if spf[i] == i:
            for j in range(i*2, N+1, i):
                if spf[j] == j:
                    spf[j] = i

    S = [0]*N
    for a in A:
        s = set()
        while a != 1:
            p = spf[a]
            if p not in s:
                S[p] += 1
            s.add(p)
            a //= p
        #print(S[0:10])


    #print(S[0:10])
    flag = False
    for i in range(len(S)):
        if S[i] == n:
            print('not coprime')
            exit()
        elif S[i] >= 2:
            flag = True
        else:
            pass
    if flag:
        print('setwise coprime')
    else:
        print('pairwise coprime')

if __name__ == '__main__':
    main()
Buggy python submission (verdict: Time Limit Exceeded):
n = int(input())

def gcd(a, b):
    while b: a, b = b, a % b
    return a
def isPrimeMR(n):
    d = n - 1
    d = d // (d & -d)
    L = [2]
    for a in L:
        t = d
        y = pow(a, t, n)
        if y == 1: continue
        while y != n - 1:
            y = (y * y) % n
            if y == 1 or t == n - 1: return 0
            t <<= 1
    return 1
def findFactorRho(n):
    m = 1 << n.bit_length() // 8
    for c in range(1, 99):
        f = lambda x: (x * x + c) % n
        y, r, q, g = 2, 1, 1, 1
        while g == 1:
            x = y
            for i in range(r):
                y = f(y)
            k = 0
            while k < r and g == 1:
                ys = y
                for i in range(min(m, r - k)):
                    y = f(y)
                    q = q * abs(x - y) % n
                g = gcd(q, n)
                k += m
            r <<= 1
        if g == n:
            g = 1
            while g == 1:
                ys = f(ys)
                g = gcd(abs(x - ys), n)
        if g < n:
            if isPrimeMR(g): return g
            elif isPrimeMR(n // g): return n // g
            return findFactorRho(g)
def primeFactor(n):
    i = 2
    ret = {}
    rhoFlg = 0
    while i*i <= n:
        k = 0
        while n % i == 0:
            n //= i
            k += 1
        if k: ret[i] = k
        i += 1 + i % 2
        if i == 101 and n >= 2 ** 20:
            while n > 1:
                if isPrimeMR(n):
                    ret[n], n = 1, 1
                else:
                    rhoFlg = 1
                    j = findFactorRho(n)
                    k = 0
                    while n % j == 0:
                        n //= j
                        k += 1
                    ret[j] = k

    if n > 1: ret[n] = 1
    if rhoFlg: ret = {x: ret[x] for x in sorted(ret)}
    return ret
ans = 0
for k in range(1, n+1):
    d = primeFactor(k)
    temp = 1
    for _, v in d.items():
        temp *= (v+1)
    ans += k*temp
print(ans)


Fixed submission:
n = int(input())
ans = 0
for i in range(1, n+1):
    m = n//i
    ans += (m*(m+1)//2)*i
print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
n,k = map(int,input().split())
A = list(map(int, input().split()))

from itertools import accumulate
for i in range(k):
    imos = [0]*(n+1)
    for j in range(n):
        d = A[j]
        l = max(0, j-d)
        r = min(n-1, j+d)
        imos[l] += 1
        imos[r+1] -= 1
    imos = list(accumulate(imos))
    for j in range(n):
        A[j] = imos[j]
    if A == [n]*n:
        break

print(*A)


Fixed submission:
n,k = map(int,input().split())
A = list(map(int, input().split()))

from itertools import accumulate
for i in range(k):
    imos = [0]*(n+1)
    for j in range(n):
        d = A[j]
        l = max(0, j-d)
        r = min(n-1, j+d)
        imos[l] += 1
        imos[r+1] -= 1
    imos = list(accumulate(imos))
    for j in range(n):
        A[j] = imos[j]
    if A == [n]*n:
        break

print(*A)
Buggy python submission (verdict: Runtime Error):
a, b, h, m = map(int, input().split())
import math
vx = 2*math.pi/(12*60)
vy = 2*math.pi/60

x = vx*(60*h+m)
y = vy*m

d = abs(x-y)
z = a**2+b**2-2*a*b*math.cos(d)
z = math.sqrt(z)
print(zf)


Fixed submission:
a, b, h, m = map(int, input().split())
import math
vx = 2*math.pi/(12*60)
vy = 2*math.pi/60

x = vx*(60*h+m)
y = vy*m

d = abs(x-y)
z = a**2+b**2-2*a*b*math.cos(d)
z = math.sqrt(z)
print(z)
Buggy python submission (verdict: Time Limit Exceeded):
def main():

    n, k = map(int, input().split())
    A = list(map(int, input().split()))
    A = [a-1 for a in A]

    L = 65
    D = [[0]*n for _ in range(L+1)]

    for i, a in enumerate(A):
        D[0][i] = a

    for i in range(L):
        for j in range(n):
            D[i+1][j] = D[i][D[i][j]]

    ans = 0
    for i in range(L):
        if k & (1 << i):
            ans = D[i][ans]
    print(ans+1)

if __name__ == '__main__':
    main()

Fixed submission:
n, k = map(int, input().split())
A = list(map(int, input().split()))

A = [a-1 for a in A]

g = [[] for i in range(n)]
for i, a in enumerate(A):
    g[i].append(a)

s = []
visit = [-1]*n
s.append(0)
visit[0] = 0
while s:
    v = s.pop()
    for u in g[v]:
        if visit[u] == -1:
            visit[u] = visit[v]+1
            s.append(u)
        else:
            temp = visit[u]
#print(temp)
#print(visit)
d = max(visit)-temp+1
#print(d)
if k < temp:
    for i in range(n):
        if visit[i] == k:
            print(i+1)
            exit()
else:
    q, r = divmod(k-temp, d)
    for i in range(n):
        if visit[i]-temp == r:
            print(i+1)
            exit()

Buggy python submission (verdict: Time Limit Exceeded):
n, k = map(int, input().split())
A = list(map(int, input().split()))
A = [a-1 for a in A]

L = 65
D = [[0]*n for _ in range(L+1)]

for i, a in enumerate(A):
    D[0][i] = a

for i in range(L):
    for j in range(n):
        D[i+1][j] = D[i][D[i][j]]

ans = 0
for i in range(L):
    if k & (1 << i):
        ans = D[i][ans]
print(ans+1)


Fixed submission:
n, k = map(int, input().split())
A = list(map(int, input().split()))

A = [a-1 for a in A]

g = [[] for i in range(n)]
for i, a in enumerate(A):
    g[i].append(a)

s = []
visit = [-1]*n
s.append(0)
visit[0] = 0
while s:
    v = s.pop()
    for u in g[v]:
        if visit[u] == -1:
            visit[u] = visit[v]+1
            s.append(u)
        else:
            temp = visit[u]
#print(temp)
#print(visit)
d = max(visit)-temp+1
#print(d)
if k < temp:
    for i in range(n):
        if visit[i] == k:
            print(i+1)
            exit()
else:
    q, r = divmod(k-temp, d)
    for i in range(n):
        if visit[i]-temp == r:
            print(i+1)
            exit()
Buggy python submission (verdict: Time Limit Exceeded):
n, m, k = map(int, input().split())
mod =998244353

if n == 1:
    print(m)
    exit()

X = [[0]*(k+2) for _ in range(k+2)]
for i in range(k+2):
    if i == 0:
        continue
    X[i][i-1] = 1
    X[i][i] = m-1

def mul(a, b):
    c = [[0]*len(b[0]) for i in range(len(a))]
    for i in range(len(a)):
        for k in range(len(b)):
            for j in range(len(b[0])):
                c[i][j] = (c[i][j] + a[i][k]*b[k][j])%mod

    return c

#A**n
def pow(a, n):
    b = [[0]*len(a) for i in range(len(a))]
    for i in range(len(a)):
        b[i][i] = 1
    while n > 0:
        if n & 1 == 1:
            b = mul(a, b)
        a = mul(a, a)
        n = n>>1
    return b

#print(X)
Y =  pow(X, n-1)
#Z0 = [0]*(k+2)
#Z[1] = m
ans = 0
#print(Y)
for i in range(k+2):
    ans += Y[i][1]*m
print(ans%mod)


Fixed submission:
n, m, k = map(int, input().split())
mod =998244353

def cmb1(n, r, mod):
    if ( r<0 or r>n ):
        return 0
    r = min(r, n-r)
    return g1[n] * g2[r] * g2[n-r] % mod

mod = 998244353 #出力の制限
N = 10**6
g1 = [1, 1] # 元テーブル
g2 = [1, 1] #逆元テーブル
inverse = [0, 1] #逆元テーブル計算用テーブル

for i in range( 2, N + 1 ):
    g1.append( ( g1[-1] * i ) % mod )
    inverse.append( ( -inverse[mod % i] * (mod//i) ) % mod )
    g2.append( (g2[-1] * inverse[-1]) % mod )

def power(a, n, mod):
  bi=str(format(n,"b")) #2進数
  res=1
  for i in range(len(bi)):
    res=(res*res) %mod
    if bi[i]=="1":
      res=(res*a) %mod
  return res

ans = 0
for x in range(k+1):
    temp = cmb1(n-1, x, mod)
    temp *= m*power(m-1, n-1-x, mod)
    ans += temp
print(ans%mod)