Python Optimisation for Round 867 G1. Magic Triples (Easy Version)
Difference between en2 and en3, changed 167 character(s)
I was doing G1 ([https://mirror.codeforces.com/contest/1822/problem/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](https://mirror.codeforces.com/contest/1822/submission/204488997).↵

<spoiler summary="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()↵
~~~~~↵
</spoiler>↵

I went to further optimise my code as shown here, [which TLEs in testcase 16](https://mirror.codeforces.com/contest/1822/submission/204528184):↵

<spoiler summary="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()↵
~~~~~↵
</spoiler>↵

Main changes:↵

1. Use faster output↵
2. Use Counter instead of array to count, which is faster (this helped the most, and got me past testcase 13)↵
3. Got rid of unnecessary variables, if else checks and array / counter access.↵
4. 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!**](https://mirror.codeforces.com/contest/1822/submission/204529110)↵

<spoiler summary="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()↵
~~~~~↵
</spoiler>↵

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:↵

1. Why does sorting the input help get rid of the Counter hash hack testcase TLE?↵
2. 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](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!

History

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