Codewars堆立方体问题求助：现有Python解法速度过慢需优化

问题背景 #

给定总容积m，需返回满足 (1^3 + 2^3 + \dots + n^3 = m) 的整数n，不存在则返回-1。示例：

• find_nb(1071225) → 45
• find_nb(91716553919377) → -1

原代码的性能问题 #

你当前的解法每次循环都重新计算从1到n的立方和，生成列表再求和的操作时间复杂度为O(n²)，当m极大时（比如测试用例135440716410000），会产生大量重复计算，导致运行速度极慢。

优化方案：利用数学公式直接求解 #

立方和有一个经典的数学公式可以彻底避免循环累加：
$$1^3 + 2^3 + \dots + n^3 = \left( \frac{n(n+1)}{2} \right)^2$$
这个公式说明，立方和必然是某个整数的平方。我们可以基于这个公式反向推导n：

1. 首先验证m是否为完全平方数：计算m的整数平方根s，若(s^2 \neq m)，直接返回-1。
2. 接下来解方程 (\frac{n(n+1)}{2} = s)，整理为二次方程 (n^2 + n - 2s = 0)，用求根公式可得：
   $$n = \frac{-1 + \sqrt{1 + 8s}}{2}$$
3. 验证这个结果是否为正整数：
   • 判别式(1+8s)必须是完全平方数
   • 计算出的n必须是整数且大于0

优化后的代码 #

import math

def find_nb(m):
    # 验证m是否是完全平方数
    s = math.isqrt(m)
    if s * s != m:
        return -1
    
    # 计算二次方程判别式并验证是否为完全平方数
    discriminant = 1 + 8 * s
    sqrt_d = math.isqrt(discriminant)
    if sqrt_d * sqrt_d != discriminant:
        return -1
    
    # 验证n是否为正整数
    if (sqrt_d - 1) % 2 != 0:
        return -1
    n = (sqrt_d - 1) // 2
    
    return n if n > 0 else -1

print(find_nb(135440716410000))  # 输出4800

代码说明 #

• 使用math.isqrt计算整数平方根，避免浮点数精度误差（Python3.8及以上版本支持）。
• 每一步都做有效性验证，快速排除不可能的情况，时间复杂度为O(1)，无论m多大都能瞬间得到结果。

内容的提问来源于stack exchange，提问作者Jensen Benny