Python关于函数的问题有人能看看嘛

编写函数，把一个自然数分解成最多4个平方数的和，然后调用这个函数。例如7604=2704+4900=52²+70²。

import itertools
n=eval(input('enter a number:')) #输入一个数
a=range(1,int(n**0.5)+1) #构造一组数，从1到n的开方，作为备选答案
haveAnswer=False #已经找到解
for i in range(2,5): # 循环，i表示从a里取出几个数来
    b=itertools.combinations_with_replacement(a,r=i) #有放回的组合（可重复选择）
    for b1 in b:
        c=sum([x**2 for x in b1]) #求平方和
        if c==n:
            haveAnswer=True
            print('{}={}'.format(n,' * '.join([str(x)+'^2' for x in b1])))
            break
    if haveAnswer:
        break

给个例子参考：

from math import sqrt
def fun(n):
    for a in range(int(sqrt(n))):
        for b in range(int(sqrt(n))):
            if a**2+b**2>n:
                break
            for c in range(int(sqrt(n))):
                if a**2+b**2+c**2>n:
                    break
                for d in range(int(sqrt(n))):
                    if a**2+b**2+c**2+d**2==n:
                        return [a,b,c,d]
n = 7604
result = fun(n)
lst = [r for r in result if r>0]
print(lst)

def cal(n):
    for i in range(n+1):
        for j in range(n+1):
            for k in range(n+1):
                for l in range(n+1):
                    if i*i+j*j+k*k+l*l == n:
                        return i,j,k,l

n = int(input())
a,b,c,d=cal(n)
print("{}={}+{}+{}+{}={}+{}+{}+{}".format(n,a**2,b**2,c**2,d**2,a,b,c,d))

from itertools import combinations

def sqrtsum(num):
    res = []
    lst = [i**2 for i in range(1, int(num**0.5)+1)]
    for i in range(1,5):
        for j in combinations(lst,i):
            if sum(j) == num:
                for k in j:
                    res.append(int(k**0.5))
                return res
    return res