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