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

№	Пользователь	Рейтинг
1	tourist	3993
2	jiangly	3743
3	orzdevinwang	3707
4	Radewoosh	3627
5	jqdai0815	3620
6	Benq	3564
7	Kevin114514	3443
8	ksun48	3434
9	Rewinding	3397
10	Um_nik	3396

№	Пользователь	Вклад
1	cry	167
2	Um_nik	163
3	maomao90	162
3	atcoder_official	162
5	adamant	159
6	-is-this-fft-	158
7	awoo	155
8	TheScrasse	154
9	Dominater069	153
10	nor	152

I was doing G1 (https://codeforces.me/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? (Edit: I found a submission that works in python https://codeforces.me/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.	Кто	Когда	Δ	Комментарий
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.

Язык

Кто

Когда

Комментарий

en5