题目解析 小R的雪球滚落计算 ｜ 豆包MarsCode AI 刷题

小R的雪球滚落计算

题目链接

问题描述

在一座高度为 H 的山上，每个高度 i 处生成了 a_i 个雪球。当雪球从海拔高度 i 滚到地面时，它的体积会膨胀 $x^i$ 倍。也就是说，雪球的初始体积为 1，滚动距离 i 会使体积变成 1 * x^i。我们需要计算所有滚落到地面的雪球的总体积，并对结果取模 $`10^9 + 7`$。

你的任务是帮助计算所有雪球滚落到地面的总体积。

---

测试样例

样例1：

输入：H = 4, x = 5, a = [1, 3, 2, 4]
输出：2830

样例2：

输入：H = 2, x = 5, a = [1, 2]
输出：55

样例3：

输入：H = 3, x = 3, a = [2, 1, 1]
输出：42

算法分析

通过循环遍历山的每一个高度 i（从 0 到 H - 1），首先调用 powerMod 方法来计算在当前高度 i 单个雪球滚落到地面时膨胀后的体积（即 x^(i + 1)，注意这里指数是 i + 1，因为高度从 0 开始计数）。

然后，将当前高度 i 处所有雪球（数量为 a[i] 个）的体积累加到 result 变量中。这里在累加前会先对单个雪球的体积乘以当前高度的雪球数量的结果进行取模操作（防止中间结果溢出），并且每累加一次也会对 result 再次取模，进一步确保整个累加过程中结果不会超出设定的范围（MOD）。

最后，将最终的总体积结果（以 long 类型计算完成后）转换为 int 类型返回，因为题目要求的返回值类型是 int，并且在之前的计算过程中已经通过取模保证了结果在 int 可表示的合理范围内（不会出现因取整导致结果错误的情况，因为取模后的数值是在合适区间的）。

这里有个很重要的算法：快速幂算法

• 初始化 result 为 1，它将作为最终结果的累积变量。
• 在循环中，每次判断指数 exp 的最低位（通过 (exp & 1) 操作来判断是否为奇数，若为 1 则最低位是 1，即指数为奇数）。如果 exp 是奇数，就将当前的 result 乘以 base，并对 mod 取模，这相当于把对应二进制位上的幂运算贡献累加到结果中。
• 然后，将 base 自身平方（即 base = base * base % mod），同时将指数 exp 右移一位（相当于 exp = exp / 2），这样就可以处理指数的下一位二进制数了。
• 不断重复这个过程，直到指数 exp 变为 0，此时 result 中存储的就是 (base^exp) % mod 的值，最后将其返回。

代码实现

public class Main {
    private static final int MOD = 1000000007;

    public static int solution(int H, int x, int[] a) {
        long result = 0;
        for (int i = 0; i < H; i++) {
            // 计算单个雪球从高度 i 滚到地面时的体积，即 x^i
            long volume = powerMod(x, i + 1, MOD);
            // 将该高度的所有雪球的体积加起来
            result += (long)a[i] * volume % MOD;
            result %= MOD; //确保结果不会溢出
        }
        return (int)result;
    }

    // 快速幂算法，用于计算 (base^exp) % mod
    private static long powerMod(long base, long exp, long mod) {
        long result = 1;
        while (exp > 0) {
            if ((exp & 1) == 1) { // 如果 exp 是奇数
                result = result * base % mod;
            }
            base = base * base % mod;
            exp >>= 1; // 相当于 exp = exp / 2
        }
        return result;
    }

    public static void main(String[] args) {
        System.out.println(solution(4, 5, new int[]{1, 3, 2, 4}) == 2830);
        System.out.println(solution(2, 5, new int[]{1, 2}) == 55);
        System.out.println(solution(3, 3, new int[]{2, 1, 1}) == 42);
    }
}

复杂度

整个代码中 solution 方法时间复杂度为 $O(n logn)$，空间复杂度为 $O(l)$

在 solution 方法中，存在一个 for 循环，循环次数取决于变量 H，循环会执行 H 次。在每次循环内部，都会调用 powerMod 方法来计算 x 的幂次。

对于 powerMod 方法，这里的 i + 1 是每次传入的指数（指数范围从 1 到 H）。因为 powerMod 方法利用快速幂算法，其时间复杂度与指数的二进制表示的位数相关，也就是对数级别（每次计算 x 的幂次时，循环次数取决于指数的二进制位数，最大为指数以 2 为底的对数）。

如前面所述，powerMod 方法中通过一个 while 循环来根据指数的二进制位进行计算。循环的次数取决于指数 exp 的二进制表示的位数，由于每次循环都会将指数 exp 右移一位（相当于除以 2），所以循环次数最多就是指数 exp 的二进制表示的位数。

总结来说，整个代码中 solution 方法时间复杂度为 $O(n logn)$，空间复杂度为 $O(l)$；其中powerMod 方法时间复杂度为 $O(n logn)$，空间复杂度为线性常熟时间复杂度 。