217小R的雪球滚落计算

题解

问题描述

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

思路

1. 理解雪球体积膨胀规则：

   • 每个雪球从高度 $i$ 滚落到地面时，其体积变为 $x^i$。
   • 同一高度 $i$ 处有 $a_i$ 个雪球，因此该高度的雪球总体积为 $a_i \times x^i$。

2. 计算总体积：

   • 对所有高度 $i$，计算 $a_i \times x^i$ 的总和。
   • 由于体积可能会非常大，需要在计算过程中对每一项取模 $10^9 + 7$，最终结果也取模。

3. 优化幂运算：

   • 使用 Python 内置的 pow 函数，利用其三个参数形式 pow(x, i + 1, mod)，高效计算 $x^{i+1} \mod (10^9 + 7)$。

代码实现

mod = int(1e9 + 7)

def solution(H: int, x: int, a: list) -> int:
    return sum(pow(x, i + 1, mod) * a[i] for i in range(len(a))) % mod

if __name__ == '__main__':
    print(solution(4, 5, [1, 3, 2, 4]) == 2830)
    print(solution(2, 5, [1, 2]) == 55)
    print(solution(3, 3, [2, 1, 1]) == 42)

复杂度分析

• 时间复杂度：$O(nlog(n))$，其中 $n$ 是数组 $a$ 的长度。每个雪球体积的计算需要 $O(log(n))$ 的时间复杂度。
• 空间复杂度：$O(1)$，只使用了常数空间存储结果。