I was doing 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 1import 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 2import 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.
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 worksimport 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!
