Problem p02888 (python)

Fixed submission:
from heapq import heappush, heappop
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  # randome is not available at Codeforces
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 show(*inp,end='\n'):
    if show_flg:
        print(*inp,end=end)
YN=['Yes','No']
mo=10**9+7
inf=float('inf')
#eps=10**(-10)
#ts=time.time()
#sys.setrecursionlimit(10**6)
input=lambda: sys.stdin.readline().rstrip()

show_flg=False
show_flg=True

n=I()
a=sorted(LI())[::-1]


def solve(a):
    
    s=set()
    for i in a:
        s.add(i)
    s=sorted(list(s))[::-1]
    L=1080
    #L=10000
    k=[0]*(L)
    ak=[0]
    ans=0
    for i in a:
        k[i]+=1
    for i in range(L):
        ak.append(ak[-1]+k[i])
    a1,a2,a3,a4=0,0,0,0
    un=len(s)
    for i in range(un):
        c=s[i]
        if k[c]>=3:
            a1+=k[c]*(k[c]-1)*(k[c]-2)//6
        if k[c]>=2:
            a2+=ak[c]*k[c]*(k[c]-1)//2
        for j in range(i+1,un):
            #show('i,j,s[i],s[j],mm',i,j,'|',s[i],s[j],s[i]-s[j]+1,'|',ak[s[i]-s[j]],ans)
            
            sc=s[j]
            a3+=(1 if 2*sc>c else 0)*k[c]*k[sc]*(k[sc]-1)//2
            if ak[sc]==0:
                break
            mm=c-sc+1
            if mm<0:
                continue
            elif mm>sc:
                continue # break
            else:
                a4+=k[c]*k[sc]*(ak[sc]-ak[mm])
            #show('ans+=',k[sc]*(k[sc]-1)//2+ak[sc]-ak[mm])
    
    ans=a1+a2+a3+a4
    #print(ans)
    return ans,a1,a2,a3,a4

def naive(a):
    b=sorted(a)[::-1]
    n=len(b)
    ans=0
    for i in range(n):
        for j in range(i+1,n):
            for k in range(j+1,n):
                if b[j]+b[k]>b[i]:
                    ans+=1
    return ans

print(solve(a)[0])

Buggy submission (verdict: Runtime Error):
from heapq import heappush, heappop
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  # randome is not available at Codeforces
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 show(*inp,end='\n'):
    if show_flg:
        print(*inp,end=end)
YN=['Yes','No']
mo=10**9+7
inf=float('inf')
#eps=10**(-10)
#ts=time.time()
#sys.setrecursionlimit(10**6)
input=lambda: sys.stdin.readline().rstrip()

show_flg=False
show_flg=True

n=I()
a=sorted(LI())[::-1]

s=set()
for i in a:
    s.add(i)
s=sorted(list(s))[::-1]
L=1080
#L=10000
k=[0]*(L)
ak=[0]
ans=0

for i in a:
    k[i]+=1
for i in range(L):
    ak.append(ak[-1]+k[i])

un=len(s)
for i in range(un):
    c=s[i]
    if k[c]>=3:
        ans+=k[c]*(k[c]-1)*(k[c]-2)//6
    if k[c]>=2:
        ans+=ak[c]*k[c]*(k[c]-1)//2
    for j in range(i+1,un):
        #show('i,j,s[i],s[j],mm',i,j,'|',s[i],s[j],s[i]-s[j]+1,'|',ak[s[i]-s[j]],ans)
        
        sc=s[j]
        if ak[sc]==0:
            break
        mm=c-sc+1
        if mm<0:
            continue
        elif mm>sc:
            continue # break
        else:
            ans+=k[c]*k[sc]*((1 if 2*k[sc]>k[c] 0)*(k[sc]-1)//2+ak[sc]-ak[mm])
        #show('ans+=',k[sc]*(k[sc]-1)//2+ak[sc]-ak[mm])
    
print(ans)

Buggy submission (verdict: Runtime Error):
from heapq import heappush, heappop
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  # randome is not available at Codeforces
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 show(*inp,end='\n'):
    if show_flg:
        print(*inp,end=end)
YN=['Yes','No']
mo=10**9+7
inf=float('inf')
#eps=10**(-10)
#ts=time.time()
#sys.setrecursionlimit(10**6)
input=lambda: sys.stdin.readline().rstrip()

show_flg=False
show_flg=True

n=I()
a=sorted(LI())[::-1]

s=set()
for i in a:
    s.add(i)
s=sorted(list(s))[::-1]
L=1080
#L=10000
k=[0]*(L)
ak=[0]
ans=0

for i in a:
    k[i]+=1
for i in range(L):
    ak.append(ak[-1]+k[i])

un=len(s)
for i in range(un):
    c=s[i]
    if k[c]>=3:
        ans+=k[c]*(k[c]-1)*(k[c]-2)//6
    elif k[c]==2:
        ans+=ak[c]
    for j in range(i+1,un):
        #show('i,j,s[i],s[j],mm',i,j,'|',s[i],s[j],s[i]-s[j]+1,'|',ak[s[i]-s[j]],ans)
        
        sc=s[j]
        if ak[sc]==0:
            break
        mm=c-sc+1
        if mm<0:
            continue
        elif mm>sc:
            continue # break
        else:
            ans+=k[sc]*(k[sc]-1)//2+ak[sc]-ak[mm]
        #show('ans+=',k[sc]*(k[sc]-1)//2+ak[sc]-ak[mm])
    
print(ans)