实现 sqrt(x):二分查找法和牛顿法

最近忙里偷闲,每天刷一道 LeetCode 的简单题保持手感,发现简单题虽然很容易 AC,但若去了解其所有的解法,也可学习到不少新的知识点,扩展知识的广度。

创作本文的思路来源于: LeetCode Problem 69. x 的平方根

简述题目大意(不想跳转链接,可以看这里):给定一个非负整数 x ,要求计算并返回 x 的平方根(取整)。例如,输入 4,则输出 2;输入 8,则输出 2(8 的平方根是 2.82842……,由于返回类型是整数,因此小数部分被舍去)。即给定一个 \(x\) ,我们要计算出 \(\lfloor \sqrt{x} \rfloor\)

最简单最直觉的方法自然是从 0 开始遍历,直到找到第一个其平方值大于 \(x\) 的数 \(n\) ,则 \(n-1\) 即是答案。对于任意的 \(x\) ,其取整后平方根一定在 \([0, x]\) 区间上,代码如下:

int sqrt(int x)
{
    if (x == 0)
        return x;
    int ans = 1;
    while (ans <= x / ans)
        ans++;
    return ans - 1;
}

这里需要注意的有两点:

  1. 第 6 行代码中, while 的判断条件可以避免溢出。很大概率上,你可能会写成 while (ans * ans <= x) ,这更自然、更直观,但当 ans 的值很大时, ans * ans 的结果可能会超过 int 类型的最大表示范围。举个例子,比如我们要计算 \(x\) 的取整平方根(其值为 \(n\) ,即 \(\lfloor \sqrt{x} \rfloor = n\) ),算法会将 ans 遍历到第一个平方超过 \(x\) 的值,即 \(n+1\) 后停止。如果 \(x\) 的值就是 int 类型能够表示的最大值,那么当 ans 遍历到 \(n+1\) 时,计算 ans * ans 的结果就超出了 int 类型的表示范围。
  2. 由于在 while 的循环判断中,我们用除法代替了乘法,因此 ans 便不能再从 0 开始遍历(否则会导致除零错误)。为此,我们可以在算法开始单独处理 \(x = 0\) 的情况,然后让 ans 从 1 开始遍历。

作为一道简单题,这种暴力朴素的算法自然是可以 AC 的。但其效率极低(需要遍历 \(O(\sqrt{n})\) 次),在 LeetCode 上的时间效率只能快过约 5% 的用户,使用 C++ 语言的运行时间平均要 90ms 以上。因此,本文提供了两种更加高效的算法:二分查找法和牛顿法。

1. 二分查找法

如果你在暴力求解的基础上继续思考,很大概率会想到用二分搜索求解。

没错,思考暴力求解的策略,我们在区间 \([0, x]\) 上搜索解,而搜索区间 \([0, x]\) 天然是有序的,自然可以用二分搜索代替线性搜索,以大大提高搜索效率。

更进一步的,我们还可以缩小我们的搜索区间。直觉告诉我们,对于一个非负整数 \(x\) ,其 \(\sqrt{x}\) 应该不会大于 \(x / 2\) (例如, \(\sqrt{25} = 5\) ,小于 \(25 / 2 = 12.5\) )。我们可以证明:

\[ \begin{aligned} &\text{设 } y = \frac{x}{2} – \sqrt{x},\text{ 则 } y^\prime = \frac{1}{2} – \frac{1}{2\sqrt{x}}, \\[2ex] &\text{令 } y^\prime = 0, \text{ 可得驻点 } x = 1, \\[2ex] &\text{当 } x > 1 \text{ 时}, y^\prime > 0, \text{ 即当 } x > 1 \text{ 时 }, y = \frac{x}{2} – \sqrt{x} \text{ 的值单调递增}, \\[2ex] &\text{可推出, 当 } x > 1 \text{ 时}, \lfloor \frac{x}{2} \rfloor – \lfloor \sqrt{x} \rfloor \text{ 的值单调递增}, \\[2ex] &\text{又因为当 } x = 2 \text{ 时}, \lfloor \frac{x}{2} \rfloor – \lfloor \sqrt{x} \rfloor = 0, \\[2ex] &\text{所以当 } x \geq 2 \text{ 时}, \lfloor \frac{x}{2} \rfloor – \lfloor \sqrt{x} \rfloor \geq 0, \text{ 即 } x \geq 2 \text{ 时},\lfloor \frac{x}{2} \rfloor \geq \lfloor \sqrt{x} \rfloor &\text{(证毕)} \end{aligned} \]

通过证明我们可以得知,当所求的 \(x\) 大于等于 \(2\) 时, sqrt(x) 的搜索空间为 \([1, x / 2]\) ,对于 \(x < 2\) 的情况,我们只要特殊处理即可(这里我们也可以得到结论:当 \(x \geq 1\) 时, \(\lfloor \frac{x}{2} \rfloor + 1 \geq \lfloor \sqrt{x} \rfloor\) ,然后单独处理 \(x < 1\) 的情况)。代码:

int sqrt(int x)
{
    if (x < 2)  // 处理特殊情况
        return x;
    
    int left = 1, right = x / 2;
    while (left > 1);
        if (mid == x / mid)
            return mid;
        else if (mid > x / mid)
            right = mid - 1;
        else
            left = mid + 1;
    }
    return right;
}

在这里解释一下最后的返回值为什么是 right 。对于二分查找,其搜索空间会不断收缩到 left == right (关于二分查找,这里不多赘述,自行手动模拟即可),此时 midleftright 三者的值相等( mid = (left + right) / 2 )。根据二分查找的搜索范围的收缩条件可知, left (或 mid )左侧的值都小于等于 \(\lfloor \sqrt{x} \rfloor\)right (或 mid )右侧的值都大于 \(\lfloor \sqrt{x} \rfloor\) ,此时( while 的最后一次循环),判断 mid 的平方与 x 的大小,有三种情况:

  1. mid == x / mid 。则在循环内直接返回 mid 的值。
  2. mid > x / mid 。这种情况下,因为 mid 左侧的值都小于等于 \(\lfloor \sqrt{x} \rfloor\) ,而 mid 的值大于 \(x\) ,则 mid-1 即是答案。而按照分支条件,执行 right = mid - 1 ,可知 right 的值正是应返回的值。此时,循环结束,应返回 right
  3. mid <= x / mid 。这种情况下, midleftright 正是计算答案(右边的值都大于 \(\lfloor \sqrt{x} \rfloor\) )。按照分支条件,执行 left = mid + 1 ,循环结束。此时, midright 的值为计算结果。

综合上面三点可知,如果 while 循环结束,则 right 保存的值一定是计算结果。

和之前的暴力法相比,使用二分查找的思想求解 sqrt(x) ,只需要循环遍历 \(O(\lg{\frac{x}{2}})\) 次;空间复杂度为 \(O(1)\)

2. 牛顿—拉弗森迭代法

牛顿—拉弗森迭代法(简称牛顿法)使用以直代曲的思想,是一种求解函数的方法,不仅仅适用于求解开方计算。当然使用牛顿法求解函数也存在很多坑,但对于求解开方而言,牛顿法是安全的。关于这一方法,你需要掌握一定的高等数学知识,想了解详细的内容,可以参考链接: 如何通俗易懂地讲解牛顿迭代法求开方?数值分析?—马同学的回答

简单的理解,可以参考图片:

图片来源: 牛顿法与拟牛顿法

给定任意一个非负整数 \(n\) ,我们想要找到一个 \(x = \lfloor \sqrt{n} \rfloor\) ,这相当于我们要计算函数 \(f(x) = x^2 – n\) 的根。我们首先需要先给出一个猜测值 \(x_0\) ,不妨令 \(x_0 = \frac{x}{2} + 1\) (证明见第一小节),然后在 \(f(x_0)\) 处作函数的切线,切线与 \(x\) 轴的交点,即为一次迭代后的值 \(x_1\) 。若 \(x_1\) 不是要得到的结果,则继续迭代,在 \(f(x_1)\) 处作函数的切线,切线与 \(x\) 轴的交点,即为第二次迭代后的值 \(x_2\) 。以此类推,直到得到 \(x_n = \lfloor \sqrt{n} \rfloor\)

现在我们来推导迭代式。对于 \(x_i\) ,其函数值为 \(f(x_i)\) ,则对于点 \((x_i, f(x_i))\) ,可得其切线方程:

\[ \begin{align} &y – f(x_i) = f(x_i)^\prime(x – x_i) \\[2ex] \implies &y – (x_i^2 – n) = 2x_i(x – x_i) \\[2ex] \implies &y + x_i^2 + n = 2x_ix \end{align} \]

又因为 \(x_{i+1}\) 为切线与 \(x\) 轴的交点,所以令 \(y=0\) ,可得:

\[ x_{i+1} = (x_i + n / x_i) / 2 \]

现在,我们就可以根据迭代式编写代码了:

int sqrt(int x)
{
    // 避免除零错误,单独处理 x = 0 的情况
    if (x == 0)
        return x;
    int t = x / 2 + 1;
    while (t > x / t)
        t = (t + x / t) / 2;
    return t;
}

为了确保算法是正确的,我们还需要一些额外的证明。首先,证明迭代式是单调递减的:

\[ x_{i+1} – x_i = \left\lfloor \frac{1}{2} (x_i + \frac{n}{x_i}) \right\rfloor – x_i = \left\lfloor \frac{1}{2} (\frac{n}{x_i} – x_i) \right\rfloor \]

可知,在区间 \([\sqrt{x}, +\infty)\) 上, \(x_{i+1} – x_i < 0\)

然后,我们还要证明迭代式是可以收敛到 \(\lfloor \sqrt{n} \rfloor\) 的:

\[ x_{i+1} = \left\lfloor \frac{1}{2} \left( x_i + \left\lfloor \frac{n}{x_i} \right\rfloor \right) \right\rfloor = \left \lfloor \frac{1}{2} (x_i + \frac{n}{x_i}) \right \rfloor \geq \left \lfloor \frac{1}{2} \times 2 \times \sqrt{x_i \cdot \frac{n}{x_i}} \right \rfloor = \lfloor \sqrt{n} \rfloor \]

因此,当 while 循环结束时,我们可以得到正确的答案。

关于牛顿法求 sqrt(x) 的时间复杂度,笔者目前也没有搞清楚,有了解的童鞋欢迎交流~。不过通过查询资料,以及实际测试,可知牛顿法的时间效率优于二分搜索法。