ManavSinghal157/NoFunEval huggingface dataset:
Task: rewrite the given python program to improve its speed of execution.

Original program (CPU time: 555 ms):
def cmb(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

A=2*10**5
mod = 10**9+7
g1 = [1, 1]
g2 = [1, 1]
inverse = [0, 1]

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

N,M,K=list(map(int,input().split()))
keisuu=cmb(N*M-2,K-2,mod)

sum_=0
for i in range(N):
    a=min(abs(i),abs(N-i-1))
    b=max(abs(i),abs(N-i-1))
    sum_+=(M**2)*((a*(a+1)//2)+(b*(b+1)//2))
for i in range(M):
    a=min(abs(i),abs(M-i-1))
    b=max(abs(i),abs(M-i-1))
    sum_+=(N**2)*((a*(a+1)//2)+(b*(b+1)//2))
print(((keisuu * (sum_//2))%mod))

Optimized program (CPU time: 399 ms):
def power_mod(a,b,mod=10**9+7):
    i,temp,box=0,b,[]
    while(2**i<=b):
        i+=1
    for j in range(i-1,-1,-1):
        box=[[j,temp//2**j]]+box
        temp-=2**j*(temp//2**j)
    box[0].append(a)
    ans=box[0][1]*a%mod
    for j in range(1,i):
        box[j].append(box[j-1][2]**2%mod)
        if box[j][1]==1:
            ans=(ans*box[j][2])%mod
    return ans
def n_func(n,mod=10**9+7):
    ans=1
    for i in range(1,n+1):
        ans=(ans*i)%mod
    return ans
def nPr(n,r,mod=10**9+7):
    ans=n_func(n-r,mod)
    ans=power_mod(ans,mod-2,mod)
    return ans*n_func(n,mod)%mod
def nCr(n,r,mod=10**9+7):
    ans=n_func(n-r,mod)*n_func(r,mod)%mod
    ans=power_mod(ans,mod-2,mod)
    return ans*n_func(n,mod)%mod

N,M,K=list(map(int,input().split()))
keisuu=nCr(N*M-2,K-2)
mod=10**9+7

sum_=0
for i in range(N):
    a=min(abs(i),abs(N-i-1))
    b=max(abs(i),abs(N-i-1))
    sum_+=(M**2)*((a*(a+1)//2)+(b*(b+1)//2))
for i in range(M):
    a=min(abs(i),abs(M-i-1))
    b=max(abs(i),abs(M-i-1))
    sum_+=(N**2)*((a*(a+1)//2)+(b*(b+1)//2))
print(((keisuu * (sum_//2))%mod))
Task: rewrite the given python program to improve its speed of execution.

Original program (CPU time: 443 ms):

import numpy as np
N = int(eval(input()))

A_ls = np.array(list(map(int,input().split(" "))))
B_ls = np.array(list(map(int,input().split(" "))))

S = 0

for i,b in enumerate(B_ls):
    if b > A_ls[i]:
        b -= A_ls[i]
        S += A_ls[i]
        if A_ls[i+1] >= b:
            A_ls[i+1] -=b
            S+=b
        else:
            S+=A_ls[i+1]
            A_ls[i+1] = 0 
    else:
        S += b

print(S)


Optimized program (CPU time: 280 ms):

N = int(eval(input()))

A_ls = list(map(int,input().split(" ")))
B_ls = list(map(int,input().split(" ")))

S = 0

for i,b in enumerate(B_ls):
    if b > A_ls[i]:
        b -= A_ls[i]
        S += A_ls[i]
        if A_ls[i+1] >= b:
            A_ls[i+1] -=b
            S+=b
        else:
            S+=A_ls[i+1]
            A_ls[i+1] = 0 
    else:
        S += b

print(S)
Task: rewrite the given python program to improve its speed of execution.

Original program (CPU time: 1095 ms):
import sys
read = sys.stdin.read
readline = sys.stdin.readline
readlines = sys.stdin.readlines

fac = [1, 1]     # 元テーブル
f_inv = [1, 1]   # 逆元テーブル
inv = [0, 1]     # 逆元テーブル計算用テーブル

def prepare(n, mod):
    for i in range(2, n+1):
        fac.append((fac[-1] * i) % mod)
        inv.append((-inv[mod % i] * (mod//i)) % mod)
        f_inv.append((f_inv[-1] * inv[-1]) % mod)

def cmb(n, r, mod):
    if n < 0 or r < 0:
        return 0
    if r > n:
        return 0

    return fac[n] * f_inv[r] * f_inv[n-r] % mod


def main():
    N,M,K = list(map(int, readline().split()))

    MOD = 998244353
    prepare(N + 10, MOD)

    if M == 1 and K == N - 1:
        print((1))
        exit()

    m = M
    for _ in range(N-1):
        m *= M - 1
        m %= MOD

    ans = 0
    for k in range(K+1):
        ans += (m * cmb(N-1, k, MOD)) % MOD
        ans %= MOD

        m *= pow(M-1, MOD-2, MOD)
        m %= MOD

    print(ans)


if __name__ == "__main__":
    main()


Optimized program (CPU time: 852 ms):
import sys
read = sys.stdin.read
readline = sys.stdin.readline
readlines = sys.stdin.readlines

# 0^0 = 1

def main():
    N,M,K = list(map(int, readline().split()))

    MOD = 998244353

    if M == 1:
        if K == N - 1:
            print((1))
        else:
            print((0))
        exit()

    m = pow(M-1, N-1, MOD)
    m_1_inv = pow(M-1, MOD-2, MOD)
    comb = 1
    ans = comb * m
    for k in range(1,K+1):
        m *= m_1_inv
        m %= MOD
        comb *= N - k
        comb %= MOD
        comb *= pow(k, MOD-2, MOD)
        
        ans += (m * comb) % MOD
        ans %= MOD

    print((ans * M % MOD))


if __name__ == "__main__":
    main()
Task: rewrite the given python program to improve its speed of execution.

Original program (CPU time: 91 ms):
import sys
sys.setrecursionlimit(10000000)

#const
dxdy=((1,0),(0,1))
#my functions here!
def pin(type=int):
    return list(map(type,input().rstrip().split()))


#your code here!
def resolve():
    S=(input())
    ans=0
    
    hidari=0
    for i,c in enumerate(S):
        if c=="W":
            ans+=i-hidari
            hidari+=1
    print(ans)

import sys
from io import StringIO
import unittest


class TestClass(unittest.TestCase):
    def assertIO(self, input, output):
        stdout, stdin = sys.stdout, sys.stdin
        sys.stdout, sys.stdin = StringIO(), StringIO(input)
        resolve()
        sys.stdout.seek(0)
        out = sys.stdout.read()[:-1]
        sys.stdout, sys.stdin = stdout, stdin
        self.assertEqual(out, output)

    def test_入力例_1(self):
        input = """BBW"""
        output = """2"""
        self.assertIO(input, output)

    def test_入力例_2(self):
        input = """BWBWBW"""
        output = """6"""
        self.assertIO(input, output)


if __name__ == "__main__":
    #unittest.main()


    resolve()#and submit 2 atcoder!



Optimized program (CPU time: 46 ms):
import sys
input= lambda: sys.stdin.readline().rstrip()
def pin(type=int):
    return list(map(type,input().split()))

#%%code

def resolve():
    
    S=(input())
    index=0
    ans=0
    for i in range(len(S)):
        if S[i]=="W":
            ans+=i-index
            index+=1
    print(ans)
#%%submit!
resolve()
Task: rewrite the given python program to improve its speed of execution.

Original program (CPU time: 170 ms):
N , L = list(map(int , input().split()))

apple_taste = []

for i in range(N):
    taste = L + i + 1 - 1
    apple_taste.append(taste)
    
import numpy as np

abs_apple_taste = np.abs(apple_taste)
m = min(abs_apple_taste)

for j in range(N):
    if m == abs_apple_taste[j]:
        apple_taste.remove(apple_taste[j])

pie_taste = 0       

for k in range(N-1):
    pie_taste += apple_taste[k]
    
print(pie_taste)

Optimized program (CPU time: 17 ms):
n,l = list(map(int,input().split()))
ans = l
eat = l

for i in range(n-1):
    l += 1
    ans += l
    
    if abs(l) < abs(eat):
        eat = l
        
print((ans-eat))
Task: rewrite the given python program to improve its speed of execution.

Original program (CPU time: 255 ms):
import sys
input = sys.stdin.buffer.readline

N, Ma, Mb = list(map(int, input().split()))
ABC = [list(map(int, input().split())) for _ in range(N)]

INF = 10 ** 15
dp = [[INF for j in range(401)] for i in range(401)]
dp[0][0] = 0

for a, b, c in ABC:
  for i in range(400, -1, -1):
    for j in range(400, -1, -1):
      if dp[i][j] != INF:
        dp[i + a][j + b] = min(dp[i + a][j + b], dp[i][j] + c)

answer = INF
for i in range(1, 401):
  for j in range(1, 401):
    if dp[i][j] != INF and i / j == Ma / Mb:
      answer = min(answer, dp[i][j])
      
print((answer if answer != INF else -1))

Optimized program (CPU time: 220 ms):
import sys
input = sys.stdin.buffer.readline

N, Ma, Mb = list(map(int, input().split()))
ABC = [list(map(int, input().split())) for _ in range(N)]

sumA = sum([ABC[i][0] for i in range(N)])
sumB = sum([ABC[i][1] for i in range(N)])

INF = 10 ** 15
dp = [[INF for j in range(sumB + 1)] for i in range(sumA + 1)]
dp[0][0] = 0

for a, b, c in ABC:
  for i in range(sumA, -1, -1):
    for j in range(sumB, -1, -1):
      if dp[i][j] != INF:
        dp[i + a][j + b] = min(dp[i + a][j + b], dp[i][j] + c)

answer = INF
for i in range(1, sumA + 1):
  for j in range(1, sumB + 1):
    if dp[i][j] != INF and i / j == Ma / Mb:
      answer = min(answer, dp[i][j])
      
print((answer if answer != INF else -1))
Task: rewrite the given python program to improve its speed of execution.

Original program (CPU time: 294 ms):
import numpy as np 
n = int(eval(input()))
x = list(map(int, input().split()))
x = np.array(x)
ans = float('inf')
idx = 200
for i in range(100):
  y = x 
  ans = min(ans, np.sum((y-i)**2))
  
print(ans)

Optimized program (CPU time: 21 ms):
N = int(eval(input()))
X = list(map(int, input().split()))
mx = max(X)
ans = float('inf')
for i in range(mx+1):
  tmp = 0
  for x in X:
    tmp += (x-(i+1))**2
  ans = min(ans, tmp)
print(ans)
Task: rewrite the given python program to improve its speed of execution.

Original program (CPU time: 760 ms):
import copy

n=int(eval(input()))
org_set=set([i for i in range(1,10**5)])
edges=[]
relations={}
nodes=[]
for i in range(n-1):
    a,b=list(map(int,input().split(" ")))
    edges.append([a,b])
    relations.setdefault(a,[]).append(b)
    nodes.append([])
nodes.append([])
nodes.append([])

max_color=0
for start in range(1,n+1):
    targets=relations.get(start,[])
    color=1
    for target in targets:
        if color in nodes[start]:
            color+=1
        else:
            pass
        nodes[target].append(color)
        if color>=max_color:
            max_color=color
        color+=1

print(max_color)
for edge in edges:
    print((nodes[edge[1]][0]))

Optimized program (CPU time: 557 ms):
n=int(eval(input()))
edges=[]
relations={}
nodes=[0,0]
for i in range(n-1):
    a,b=list(map(int,input().split(" ")))
    edges.append([a,b])
    relations.setdefault(a,[]).append(b)
    nodes.append(0)

max_color=0
for start in range(1,n+1):
    targets=relations.get(start,[])
    color=1
    for target in targets:
        if color == nodes[start]:
            color+=1
        else:
            pass
        nodes[target]=color
        if color>=max_color:
            max_color=color
        color+=1

print(max_color)
for edge in edges:
    print((nodes[edge[1]]))