Python Optimisation for Round 867 G1. Magic Triples (Easy Version)

№	Пользователь	Рейтинг
1	tourist	4009
2	jiangly	3839
3	Radewoosh	3646
4	jqdai0815	3620
4	Benq	3620
6	orzdevinwang	3612
7	Geothermal	3569
7	cnnfls_csy	3569
9	ecnerwala	3494
10	Um_nik	3396

№	Пользователь	Вклад
1	Um_nik	164
2	maomao90	160
3	-is-this-fft-	159
4	atcoder_official	158
4	cry	158
4	awoo	158
7	adamant	155
8	nor	154
9	TheScrasse	153
10	maroonrk	152

I was doing G1 (https://codeforces.net/contest/1822/problem/G1) and came up with this solution which is O(n * sqrt(m)) as stated in the editorial. This TLEs in testcase 13 due to constant factors.

TLE Code 1

import sys
 
input = sys.stdin.readline
 
 
def solve(n, arr):
    counts = [0] * (max(arr) + 1)
    for i in arr:
        counts[i] += 1
 
    ans = 0
    for i in set(arr):
        b = 1
        while i * b * b < len(counts):
            x, y, z = counts[i], counts[i * b], counts[i * b * b]
            if b == 1:
                ans += x * (x - 1) * (x - 2)
            else:
                ans += x * y * z
            b += 1
 
    print(ans)
 
 
def main():
    t = int(input())
    for _ in range(t):
        n = int(input())
        arr = list(map(int, input().split()))
        solve(n, arr)
 
 
main()

I went to further optimise my code as shown here, which TLEs in testcase 16:

TLE Code 2

import sys
from collections import Counter
 
input = sys.stdin.readline
 
 
def solve(n, arr):
    c = Counter(arr)
    m = max(arr)
 
    ans = 0
    for i, x in c.items():
        ans += x * (x - 1) * (x - 2)
        for b in range(2, 10**3 + 1):
            if i * b * b > m:
                break
            ans += x * c[i * b] * c[i * b * b]
 
    sys.stdout.write(f"{ans}\n")
 
 
def main():
    t = int(input())
    for _ in range(t):
        n = int(input())
        arr = list(map(int, input().split()))
        solve(n, arr)
 
 
main()

Main changes:

Use faster output
Use Counter instead of array to count, which is faster (this helped the most, and got me past testcase 13). In another solution, I initialised the count array at the very start instead of for each testcase which helped me pass everything.
Got rid of unnecessary variables, if else checks and array / counter access.
Use for loop instead of while loop

However, this still TLEs due to Testcase 16, which is a hash hack for Counter. If I sort the input beforehand, my code passes!

Code that works

import sys
from collections import Counter
 
input = sys.stdin.readline
 
 
def solve(arr):
    arr.sort()
    c = Counter(arr)
    m = arr[-1]
 
    ans = 0
    for i, x in c.items():
        ans += x * (x - 1) * (x - 2)
        for b in range(2, 10**3 + 1):
            if i * b * b > m:
                break
            ans += x * c[i * b] * c[i * b * b]
 
    sys.stdout.write(f"{ans}\n")
 
 
def main():
    t = int(input())
    for _ in range(t):
        input()
        arr = list(map(int, input().split()))
        solve(arr)
 
 
main()

Note that Python is too slow to pass all testcases (it fails at testcase 13), but PyPy works.

I have a few questions for the people well versed in python:

Why does sorting the input help get rid of the Counter hash hack testcase TLE?
Is it possible to make Python pass all testcases? (Edit: I found a submission that works in python https://codeforces.net/contest/1822/submission/203361435)

I hope you can answer my questions, if not, I hope you have learnt something from my blog. Thank you!