Improve solution

Author: MaximSmolskiy

project_euler/problem_012/sol1.py

@@ -21,17 +21,20 @@
 What is the value of the first triangle number to have over five hundred
 divisors?
 """
-from math import sqrt
 
 
 def count_divisors(n):
-    nDivisors = 0
-    for i in range(1, int(sqrt(n)) + 1):
-        if n % i == 0:
-            nDivisors += 2
-    # check if n is perfect square
-    if n ** 0.5 == int(n ** 0.5):
-        nDivisors -= 1
+    nDivisors = 1
+    i = 2
+    while i * i <= n:
+        multiplicity = 0
+        while n % i == 0:
+            n //= i
+            multiplicity += 1
+        nDivisors *= multiplicity + 1
+        i += 1
+    if n > 1:
+        nDivisors *= 2
     return nDivisors