python解决简单的数学问题

问题：(5·3中考数学P5-14)
定义：如果一个数k可以表示为2个正整数(m,n)的平方之差(m-n>1)，则称k为“智慧优数”
将智慧优数从小到大排列，求第3个智慧优数和第23个智慧优数。

解答(numba加速计算)

from numba import njit

@njit
def getNums(mMax:int):
    """ 获取所有m<=mMax的智慧优数 """
    result = []
    for m in range(1,mMax+1):
        for n in range(1,mMax+1):
            if (m - n)>1:
                result.append(m**2-n**2)

    return result


@njit
def sort(nums: list[int]):
    """ 排序 """
    for i in range(len(nums) - 1):
        for j in range(len(nums) - 1 - i):
            if nums[j] > nums[j + 1]:
                nums[j], nums[j + 1] = nums[j + 1], nums[j]
    return nums

def main():
    nums = getNums(1000)
    return sort(list(set(nums)))


result = main()
print(len(result))
print(result[3-1])
print(result[23-1])