Problem p02824 (python)

Fixed submission:
from heapq import heappush,heappop,heapify
from collections import deque,defaultdict,Counter
import itertools
from itertools import permutations,combinations,groupby
import sys
import bisect
import string
import math
import time
import random
def S_():
    return input()
def LS():
    return [i for i in input().split()]
def I():
    return int(input())
def MI():
    return map(int,input().split())
def LI():
    return [int(i) for i in input().split()]
def LI_():
    return [int(i)-1 for i in input().split()]
def StoI():
    return [ord(i)-97 for i in input()]
def ItoS(nn):
    return chr(nn+97)
def LtoS(ls):
    return ''.join([chr(i+97) for i in ls])
def GI(V,E,Directed=False,index=0):
    org_inp=[]
    g=[[] for i in range(n)]
    for i in range(E):
        inp=LI()
        org_inp.append(inp)
        if index==0:
            inp[0]-=1
            inp[1]-=1
        if len(inp)==2:
            a,b=inp
            g[a].append(b)
            if not Directed:
                g[b].append(a)
        elif len(inp)==3:
            a,b,c=inp
            aa=(inp[0],inp[2])
            bb=(inp[1],inp[2])
            g[a].append(bb)
            if not Directed:
                g[b].append(aa)
    return g,org_inp
def bit_combination(k,n=2):
    rt=[]
    for tb in range(n**k):
        s=[tb//(n**bt)%n for bt in range(k)]
        rt+=[s]
    return rt
def show(*inp,end='\n'):
    if show_flg:
        print(*inp,end=end)
YN=['Yes','No']
mo=10**9+7
inf=float('inf')
l_alp=string.ascii_lowercase
u_alp=string.ascii_uppercase
ts=time.time()
#sys.setrecursionlimit(10**5)
input=lambda: sys.stdin.readline().rstrip()

def ran_input():
    import random
    n=random.randint(4,16)
    rmin,rmax=1,10
    a=[random.randint(rmin,rmax) for _ in range(n)]
    return n,a

show_flg=False
show_flg=True

ans=0

n,m,v,p=LI()
a=LI()
b=sorted(a)[::-1]

def check(k):
    tot=m*v
    if k<=p-1:
        return True
    if b[p-1]>b[k]+m:
        return False
    tot-=m*p
    mx=b[k]+m
    for i in range(p-1,n):
        if i==k:
            continue
        tot-=min(max(0,mx-b[i]),m)

    if tot<=0:
        return True
    else:
        return False
        
l,r=0,n

while r-l>1:
    mid=(l+r)//2
    if check(mid):
        l=mid
    else:
        r=mid

ok=b[l]
ans=sum([1 if i>=ok else 0 for i in b])
print(ans)

Buggy submission (verdict: Runtime Error):
from heapq import heappush,heappop,heapify
from collections import deque,defaultdict,Counter
import itertools
from itertools import permutations,combinations
import sys
import bisect
import string
import math
import time
import random
def I():
    return int(input())
def MI():
    return map(int,input().split())
def LI():
    return [int(i) for i in input().split()]
def LI_():
    return [int(i)-1 for i in input().split()]
def StoI():
    return [ord(i)-97 for i in input()]
def ItoS(nn):
    return chr(nn+97)
def GI(V,E,Directed=False,index=0):
    org_inp=[]
    g=[[] for i in range(n)]
    for i in range(E):
        inp=LI()
        org_inp.append(inp)
        if index==0:
            inp[0]-=1
            inp[1]-=1
        if len(inp)==2:
            a,b=inp
            g[a].append(b)
            if not Directed:
                g[b].append(a)
        elif len(inp)==3:
            a,b,c=inp
            aa=(inp[0],inp[2])
            bb=(inp[1],inp[2])
            g[a].append(bb)
            if not Directed:
                g[b].append(aa)
    return g,org_inp
def bit_combination(k,n=2):
    rt=[]
    for tb in range(n**k):
        s=[tb//(n**bt)%n for bt in range(k)]
        rt+=[s]
    return rt
def show(*inp,end='\n'):
    if show_flg:
        print(*inp,end=end)
YN=['Yes','No']
mo=10**9+7
inf=float('inf')
l_alp=string.ascii_lowercase
u_alp=string.ascii_uppercase
ts=time.time()
#sys.setrecursionlimit(10**5)
input=lambda: sys.stdin.readline().rstrip()

def ran_input():
    import random
    n=random.randint(4,16)
    rmin,rmax=1,10
    a=[random.randint(rmin,rmax) for _ in range(n)]
    return n,a

show_flg=False
#show_flg=True

ans=0

n,m,v,p=LI()
a=sorted(LI())

def check(x):
    y=min(bisect.bisect(a,a[x]+m),n-1)
    #show('y',y)
    vv=max(v-(n-y+1),0)
    if vv==0 and n-bisect.bisect(a,a[x]+m+1)<p:
        return True
    show('y,a[y],vv=',y,a[y],vv)
    if IsDistributable(m,vv,[a[x]+m - i for i in a[x+1:y][::-1]+a[:x][::-1]]):
        return True
    else:
        return False
    return y

def IsDistributable(M,V,B):
    show('M,V,B=',M,V,B)
    p=M*V
    r=0
    pre=0
    W=len(B)
    for i in range(W):
        if i!=0 and pre==B[i]:
            continue
        
        avl=(B[i]-pre)*(n-i)-r
        p-=avl+r
        r=avl%V
        if V<n-i:
            return False
        pre=B[i]
        if p<=0:return True


# initial value
ok=n-1
ng=-1

while abs(ok-ng)>1:
    mid=(ok+ng)//2
    if check(mid):
        ok=mid
    else:
        ng=mid
ans=ok
show('ans')
print(n-ans)
x=3
show('x,check=',x,check(x))
show('n,m,v,p=',n,m,v,p)
show(a)