Problem p02774 (python)

Fixed submission:
import numpy as np

N, K=map(int, input().split())
A=np.sort(list(map(int, input().split())))

p=np.sum(A>0)
q=np.sum(A<0)
Z=N-p-q
if p>0:
  AP=A[-p:]
else:
  AP=[]
if q>0:
  AN=A[:q]
else:
  AN=[]

b=(N-Z)*Z+Z*(Z-1)//2
a=p*q

if K<=a:
  x=AP[-1]*AN[0]-1
  y=AP[0]*AN[-1]
  while y-x>1:
    tmp=(y+x)//2
    counter=np.searchsorted(AN, tmp//AP, side='right').sum()
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)
elif K<=a+b:
  print(0)
else:
  x=0
  y=10**18
  K=K-a-b
  if q<=1:
    while y-x>1:
      tmp=(y+x)//2
      counter=(np.searchsorted(AP, tmp//AP, side='right').sum()-np.searchsorted(AP, int(pow(tmp, 1/2)), side='right'))//2
      if counter<K:
        x=tmp
      else:
        y=tmp
    print(y)
    exit()
    
  AN=AN[::-1]
  AN=-AN
  
  if p<=1:
    while y-x>1:
      tmp=(y+x)//2
      counter=(np.searchsorted(AN, tmp//AN, side='right').sum()-np.searchsorted(AN, int(pow(tmp, 1/2)), side='right'))//2
      if counter<K:
        x=tmp
      else:
        y=tmp
    print(y)
    exit()
    
  while y-x>1:
    tmp=(y+x)//2
    counter1=(np.searchsorted(AP, tmp//AP, side='right').sum()-np.searchsorted(AP, int(pow(tmp, 1/2)), side='right'))//2
    counter2=(np.searchsorted(AN, tmp//AN, side='right').sum()-np.searchsorted(AN, int(pow(tmp, 1/2)), side='right'))//2
    counter=counter1+counter2
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)

Buggy submission (verdict: Runtime Error):
import numpy as np
def f(nums, a):
  return 0  if nums.size<=1 else (np.searchsorted(nums, a//nums, side='right').sum()-np.searchsorted(nums, int(pow(a, 1/2)), side='right'))//2

N, K=map(int, input().split())
A=np.sort(list(map(int, input().split())))

p=np.sum(A>0)
q=np.sum(A<0)
Z=N-p-q
AP=A[-p:] if p>0 else np.array([])
AN=A[:q] if q>0 else np.array([])

b=(N-Z)*Z+Z*(Z-1)//2
a=p*q

if K<=a:
  x=AP[-1]*AN[0]-1
  y=AP[0]*AN[-1]
  while y-x>1:
    tmp=(y+x)//2
    counter=np.searchsorted(AN, tmp//AP, side='right').sum()
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)
elif K<=a+b:
  print(0)
else:
  if q>1:
    AN=-AN[::-1]
  x=0
  y=max(AN[-2]*AN[-1], AP[-2]*AP[-1])
  K=K-a-b
  while y-x>1:
    tmp=(y+x)//2
    if f(AP, tmp)+f(AN, tmp)<K:
      x=tmp
    else:
      y=tmp
  print(y)

Buggy submission (verdict: Runtime Error):
import numpy as np

N, K=map(int, input().split())
A=np.sort(list(map(int, input().split())))

p=np.sum(A>0)
q=np.sum(A<0)
Z=N-p-q
AP=A[-p:]
if q>0:
  AN=A[:q]

b=(N-Z)*Z+Z*(Z-1)//2
a=p*q

if K<=a:
  x=AP[-1]*AN[0]-1
  y=AP[0]*AN[-1]
  while y-x>1:
    tmp=(y+x)//2
    counter=np.searchsorted(AN, tmp//AP, side='right').sum()
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)
elif K<=a+b:
  print(0)
else:
  K=K-a-b
  x=0
  y=10**18
  AN=AN[::-1]
  AN=-AN
  while y-x>1:
    tmp=(y+x)//2
    counter1=(np.searchsorted(AP, tmp//AP, side='right').sum()-np.searchsorted(AP, int(pow(tmp, 1/2)), side='right'))//2
    counter2=(np.searchsorted(AN, tmp//AN, side='right').sum()-np.searchsorted(AN, int(pow(tmp, 1/2)), side='right'))//2
    counter=counter1+counter2
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)

Buggy submission (verdict: Runtime Error):
import numpy as np

N, K=map(int, input().split())
A=np.sort(list(map(int, input().split())))

p=np.sum(A>0)
q=np.sum(A<0)
Z=N-p-q
AP=A[-p:]
if q>0:
  AN=A[:q]

b=(N-Z)*Z+Z*(Z-1)//2
a=p*q

if K<=a:
  x=AP[-1]*AN[0]-1
  y=AP[0]*AN[-1]
  while y-x>1:
    tmp=(y+x)//2
    counter=np.searchsorted(AN, tmp/AP, side='right').sum()
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)
elif K<=a+b:
  print(0)
else:
  K=K-a-b
  x=0
  y=10**18
  AN=AN[::-1]
  AN=-AN
  while y-x>1:
    tmp=(y+x)//2
    counter1=(np.searchsorted(AP, tmp//AP, side='right').sum()-np.searchsorted(AP, int(pow(tmp, 1/2)), side='right'))//2
    counter2=(np.searchsorted(AN, tmp//AN, side='right').sum()-np.searchsorted(AN, int(pow(tmp, 1/2)), side='right'))//2
    counter=counter1+counter2
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)

Buggy submission (verdict: Runtime Error):
import numpy as np

N, K=map(int, input().split())
A=np.sort(list(map(int, input().split())))

p=np.sum(A>0)
q=np.sum(A<0)
Z=N-p-q
AP=A[-p:]
if q>0:
  AN=A[:q]

b=(N-Z)*Z+Z*(Z-1)//2
a=p*q

if K<=a:
  x=AP[-1]*AN[0]-1
  y=AP[0]*AN[-1]
  while y-x>1:
    tmp=(y+x)//2
    counter=np.searchsorted(AN, tmp/AP, side='right').sum()
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)
elif K<=a+b:
  print(0)
else:
  K=K-a-b
  x=0
  y=10**18
  AN=AN[::-1]
  AN=-AN
  while y-x>1:
    tmp=(y+x)//2
    counter1=(np.searchsorted(AP, tmp//AP, side='right').sum()-np.searchsorted(AP, pow(tmp, 1/2), side='right'))//2
    counter2=(np.searchsorted(AN, tmp//AN, side='right').sum()-np.searchsorted(AN, pow(tmp, 1/2), side='right'))//2
    counter=counter1+counter2
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)

Buggy submission (verdict: Runtime Error):
import numpy as np

N, K=map(int, input().split())
A=np.sort(list(map(int, input().split())))

p=np.sum(A>0)
q=np.sum(A<0)
Z=N-p-q
AP=A[-p:]
if q>0:
  AN=A[:q]

b=(N-Z)*Z+Z*(Z-1)//2
a=p*q

if K<=a:
  x=AP[-1]*AN[0]-1
  y=AP[0]*AN[-1]
  while y-x>1:
    tmp=(y+x)//2
    counter=np.searchsorted(AN, tmp/AP, side='right').sum()
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)
elif K<=a+b:
  print(0)
else:
  K=K-a-b
  x=0
  y=10**18
  AN=AN[::-1]
  AN=-AN
  while y-x>1:
    tmp=(y+x)//2
    counter1=(np.searchsorted(AP, tmp/AP, side='right').sum()-np.searchsorted(AP, pow(tmp, 1/2), side='right'))//2
    counter2=(np.searchsorted(AN, tmp/AN, side='right').sum()-np.searchsorted(AN, pow(tmp, 1/2), side='right'))//2
    counter=counter1+counter2
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)

Buggy submission (verdict: Runtime Error):
import numpy as np

N, K=map(int, input().split())
A=np.sort(list(map(int, input().split())))

p=np.sum(A>0)
q=np.sum(A<0)
Z=N-p-q
AP=A[-p:]
if q>0:
  AN=A[:q]

b=(N-Z)*Z+Z*(Z-1)//2
a=p*q

if K<=a:
  x=AP[-1]*AN[0]
  y=AP[0]*AN[-1]
  while y-x>1:
    tmp=(y+x)//2
    counter=np.searchsorted(AN, tmp/AP, side='right').sum()
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)
elif K<=a+b:
  print(0)
else:
  K=K-a-b
  x=1
  y=10**18
  AN*=-1
  AN=AN[::-1]
  while y-x>1:
    tmp=(y+x)//2
    counter1=(np.searchsorted(AP, tmp/AP, side='right').sum()-np.searchsorted(AP, tmp**(1/2), side='right'))/2
    counter2=(np.searchsorted(AN, tmp/AN, side='right').sum()-np.searchsorted(AN, tmp**(1/2), side='right'))/2
    counter=counter1+counter2
    if counter<K:
      x=tmp
    else:
      y=tmp
  print(y)

Buggy submission (verdict: Runtime Error):
import heapq
N, K=map(int, input().split())
A=list(map(int, input().split()))
A.sort()
if A[0]>=0:
  print(kantan(A))
  exit()
if A[-1]<=0:
  A=A[::-1]
  print(kantan(A))
  exit()
  
x=0
y=N-1
while y-x>1:
  if A[(x+y)//2]<0:
    x=(x+y)//2
  else:
    y=(x+y)//2
w=y
z=N-1
while z-w>1:
  if A[(w+z)//2]<1:
    w=(w+z)//2
  else:
    z=(w+z)//2
    
if A[y]!=0:
    y=z
    
po=y*(y-1)//2+(N-z)*(N-z-1)//2
ne=y*(N-z)
ze=(z-y)*(N-(z-y))+(z-y)*(z-y-1)//2

if K<=ne:
  nums=[[A[0]*A[-1], 0, -1]]
  heapq.heapify(nums)
  counter=0
  sumi=[[0]*N for _ in range(N)]
  while counter<K:
    tmp=heapq.heappop(nums)
    if A[tmp[1]+1]<0 and sumi[tmp[1]+1][tmp[2]]==0:
      heapq.heappush(nums, [A[tmp[1]+1]*A[tmp[2]], tmp[1]+1, tmp[2]])
      sumi[tmp[1]+1][tmp[2]]=1
      sumi[tmp[2]][tmp[1]+1]=1
    if A[tmp[2]-1]>0 and sumi[tmp[1]][tmp[2]+1]==0:
      heapq.heappush(nums, [A[tmp[1]]*A[tmp[2]-1], tmp[1], tmp[2]-1])
      sumi[tmp[1]][tmp[2]+1]=1
      sumi[tmp[2]+1][tmp[1]]=1
    counter+=1
  print(tmp[0])
  exit()
  
def kantan(A):
  nums=[[A[0]*A[1], 0, 1]]
  heapq.heapify(nums)
  counter=0
  sumi=[[0]*N for _ in range(N)]
  while counter<K:
    tmp=heapq.heappop(nums)
    if tmp[2]-tmp[1]>1 and sumi[tmp[1]+1][tmp[2]]==0:
      heapq.heappush(nums, [A[tmp[1]+1]*A[tmp[2]], tmp[1]+1, tmp[2]])
      sumi[tmp[1]+1][tmp[2]]=1
      sumi[tmp[2]][tmp[1]+1]=1
    if tmp[2]<N-1 and sumi[tmp[1]][tmp[2]+1]==0:
      heapq.heappush(nums, [A[tmp[1]+1]*A[tmp[2]], tmp[1], tmp[2]+1])
      sumi[tmp[1]][tmp[2]+1]=1
      sumi[tmp[2]+1][tmp[1]]=1
    counter+=1
  return tmp[0]
  
if K<=ne+ze:
  print(0)
  exit()

nums=[[A[y-2]*A[y-1], y-2, y-1, -1], [A[z]*A[z+1], z, z+1, 1]]
heapq.heapify(nums)
counter=ne+ze
sumi=[[0]*N for _ in range(N)]
while counter<K:
  tmp=heapq.heappop(nums)
  if tmp[3]>0:
    if tmp[2]-tmp[1]>1 and sumi[tmp[1]+1][tmp[2]]==0:
      heapq.heappush(nums, [A[tmp[1]+1]*A[tmp[2]], tmp[1]+1, tmp[2], 1])
      sumi[tmp[1]+1][tmp[2]]=1
      sumi[tmp[2]][tmp[1]+1]=1
    if tmp[2]<N-1 and sumi[tmp[1]][tmp[2]+1]==0:
      heapq.heappush(nums, [A[tmp[1]]*A[tmp[2]+1], tmp[1], tmp[2]+1, 1])
      sumi[tmp[1]][tmp[2]+1]=1
      sumi[tmp[2]+1][tmp[1]]=1
  else:
    if tmp[2]-tmp[1]>1 and sumi[tmp[1]][tmp[2]-1]==0:
      heapq.heappush(nums, [A[tmp[1]]*A[tmp[2]-1], tmp[1], tmp[2]-1, -1])
      sumi[tmp[1]][tmp[2]-1]=1
      sumi[tmp[2]-1][tmp[1]]=1
    if tmp[2]>0 and sumi[tmp[1]-1][tmp[2]]==0:
      heapq.heappush(nums, [A[tmp[1]]*A[tmp[2]-1], tmp[1]-1, tmp[2], -1])
      sumi[tmp[1]-1][tmp[2]]=1
      sumi[tmp[2]][tmp[1]-1]=1
  counter+=1
print(tmp[0]) #\(^o^)/ｺﾞﾘｵｼｼﾀｲｼﾞﾝｾｲﾀﾞｯﾀ