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

#	User	Rating
1	tourist	3757
2	jiangly	3647
3	Benq	3581
4	orzdevinwang	3570
5	Geothermal	3569
5	cnnfls_csy	3569
7	Radewoosh	3509
8	ecnerwala	3486
9	jqdai0815	3474
10	gyh20	3447

#	User	Contrib.
1	maomao90	171
2	adamant	164
3	awoo	163
4	TheScrasse	159
5	nor	155
6	maroonrk	154
7	-is-this-fft-	152
8	Petr	147
9	orz	146
10	pajenegod	145

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)
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?

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