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

#	User	Rating
1	tourist	3985
2	jiangly	3814
3	jqdai0815	3682
4	Benq	3529
5	orzdevinwang	3526
6	ksun48	3517
7	Radewoosh	3410
8	hos.lyric	3399
9	ecnerwala	3392
9	Um_nik	3392

#	User	Contrib.
1	cry	169
2	maomao90	162
2	Um_nik	162
4	atcoder_official	161
5	djm03178	158
6	-is-this-fft-	157
7	adamant	155
8	awoo	154
8	Dominater069	154
10	luogu_official	150

I was doing G1 (https://mirror.codeforces.com/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). Initialising the count array at the very start instead of for each testcase also helped me to 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://mirror.codeforces.com/contest/1822/submission/203361435)

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

Rev.	By	When	Δ	Comment
en5	drugkeeper	2023-05-05 11:17:28	40
en4	drugkeeper	2023-05-05 11:16:13	112
en3	drugkeeper	2023-05-05 10:40:41	167
en2	drugkeeper	2023-05-05 10:11:16	325
en1	drugkeeper	2023-05-05 10:08:17	3049	Initial revision (published)

Rev.

Lang.

When

Comment

en5