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

Question

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

Answer 1

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