How to cut down on run-time for finding dimensions of a sphere Python?

Question

I'm supposed to write code that prints the total number of integer solutions to the inequality x^2 + y^2 + z^2 <= n, where n is a user-inputted integer that is between 1 and 2000 (inclusive), and adds all the previous numbers of solutions (E.g. n=1 returns 7 and n=2 is 19, but it returns 26 and so forth). Here is my code:

import math
import itertools
n = int(input("Please enter an integer between 1 and 2000: "))


def sphereCombos(radius):
  for i in range(1, radius+1):
    count = 0
    CombosList = []
    rad = int(math.sqrt(i))
    range_for_x= range(-rad, rad + 1) 
    range_for_y= range(-rad, rad + 1)
    range_for_z= range(-rad, rad + 1)
    total_perms = list(itertools.product(range_for_x, range_for_y, range_for_z))
    
  for x, y, z in total_perms:
      if x*x+ y*y + z*z <= i:
        count = count + 1
  return count

possible_combos = 0
for i in range(1, n + 1):
  possible_combos = possible_combos + sphereCombos(i)
  print(possible_combos)
  

The code works exactly as it's supposed to, but the problem is when n is set to be 2000, the program takes way too long, and I need to get it to run in 2 minutes or less. I thought using .product() would make it much faster than using three nested for loops, but that didn't end up being super true. Is there any way for me to cut down on run time?

Answer 1

The code is working because i is in global scope. Notice 'sphereCombos(i)' passes i to radius. But the function is actually using global i in rad = int(math.sqrt(i)) and in if x * x + y * y + z * z <= i:

import math
import itertools

n = int(input("Please enter an integer between 1 and 2000: "))


def sphereCombos(radius):
    count = 0
    rad = int(math.sqrt(radius)) # Changed from i to radius
    range_for_x = range(-rad, rad + 1)
    range_for_y = range(-rad, rad + 1)
    range_for_z = range(-rad, rad + 1)
    total_perms = list(
        itertools.product(range_for_x, range_for_y, range_for_z))

    for x, y, z in total_perms:
        if x * x + y * y + z * z <= radius: # Changed from i to radius
            count = count + 1
    return count


possible_combos = 0
for i in range(1, n + 1):
    possible_combos = possible_combos + sphereCombos(i)
    print(possible_combos)