7-2 国王的奖励——分数取模、分治思想、快速幂、int64的乘法模运算

1 问题

题面描述

国际象棋是很久以前由一个印度人 Shashi 发明的，当他把该发明献给国王的时候，国王很高兴，就许诺可以给这个发明人任何他想要的奖赏。Shashi 要求以这种方式给他一些粮食：棋盘的第 1 个方格里只放 1 粒麦粒，第 2 格 q 个,第 3 格 q^2^ 个,第 4 格 q^3^ 个……，直到 n 个格子全部放满。这个奖赏最终会是什么样子的呢？Shashi 已经有算法了，请你算一算吧。

当然 Shashi 并不关心具体的数有多少了，他有一个检验你的答案是否与他心意相通的办法:把你求出的答案对 100000007 取模看看和 Shashi 算的是否一样就行了。

输入描述

本题输入具有多组样例。

第一行为一个数 t (1 ≤ t ≤ 500)，代表测试数据的组数。

之后的 t 行，每行两个数 q, n (1 ≤ q, n ≤ 10^18^)，含义见题目描述。

输出描述

输出共 t 行，每行一个数，为你算的答案。 友情提醒：你算的答案应该比 100000007 小。

样例描述

样例输入：

3
2 1
2 2
2 3

样例输出：

1
3
7

样例解释

对于样例：

• q = 2, n = 1 时仅有1个麦粒
• q = 2, n = 2 时有 1 + 2 = 3 个麦粒
• q = 2, n = 3 时有 1 + 2 + 2^2^ = 7 个麦粒

代码长度限制｜16 KB
时间限制｜1000 ms
内存限制｜256 MB

2 分析

2.1 数学抽象

该问题本质为等比数列求和问题，$1$ 为首项，$q$ 为公比，$n$ 为项数。 需要注意：

• 每次输入的 $q$, $n$ 可能不同；
• $q$, $n$ 的范围超过 int32，需要用 int64 来存储；
• 结果要对 $100000007$ 取模；
• 计算过程中涉及到大数乘法，直接用 * 运算符会造成数据溢出，需要用自己设计的大数乘法模运算作为替代。

2.2 解决方法

2.2.1 等比求和

等比数列求和方法一般有直接求和、二分递归以及运用等比求和通项公式。

对于 $q=1$ 的情况，显然等比求和通项公式是最好的选择，取模对此并无影响，时间复杂度为 $O(1)$，因此问题集中在 $q>1$ 的情况之上。

本文中我们将依次对 $q>1$ 情况下的等比求和通项公式法和二分递归法进行讨论分析。

2.2.1.1 求和公式：分数取模——乘法逆元

等比数列求和通项公式如下：

$$S_{n}=\dfrac{1-q^n}{1-q}$$

求和部分时间复杂度只有 $O(1)$，主要取决于求解 $q^n$ 的算法。

但模运算对除法不存在分配率，即

$$\dfrac{a \bmod p}{b \bmod p} \ne \dfrac{a}{b} \bmod p$$

不能直接使用通项公式。

对分数取模有两种方法：乘法逆元和变换模值，我们在此仅对逆元进行讨论，变换模值可参见 [参考文章1]。

若在 $\bmod~p$ 意义下，对于一个整数 $a$，有 $a \cdot b \equiv 1 \pmod p$，那么这个整数 $b$ 与 $a$ 互为乘法逆元

求乘法逆元的方法主要为：将 $a \cdot b \equiv 1 \pmod p$ 转化为 $ax+py=1$，用拓展欧几里得算法解二元一次方程解出 $x$。

但本题中 $p=100000007$ 是质数，因此可以用费马小定理（详见 [参考文章3]）直接求解。

费马小定理：

若 $a$ 为整数，$p$ 为质数，则有 $a^p \bmod p \equiv a$

当 $a$ 不是 $p$ 的倍数时，可写为 $a^{p-1} \bmod p \equiv 1$

故有

$$\dfrac{a}{b} \bmod p = (b^{p-1} \bmod p) \cdot (\dfrac{a}{b} \bmod p) = a \cdot b^{p-2} \bmod p$$

即

$$S_n = (1-q^n) \cdot (1-q)^{p-2} \bmod p$$

警告

若 $b$ 是 $p$ 的倍数则不能用逆元。因为此时 $b \bmod p \equiv 0$，无法求解逆元。

代码实现如下：

// q: 公比; n：项数; mod: 模; qpow: 求幂函数
if (q == 1) {
	printf("%lld\n", n % mod);
} else {
	printf("%lld\n", (qpow(q, n) - 1) * qpow(q - 1, mod - 2) % mod);
}

2.2.1.2 二分递归：数学推导——分治思想

如前文所述，当 $b$ 是 $p$ 的倍数时，无法求解逆元。 当数据中存在上述关系时，我们需要另寻他法，此处选择二分递归法进行等比求和。

要进行二分递归，我们首先要进行数学推导，得出等比数列前 $n$ 项和 $S_n$ 的递推公式。运用分治思想，我们对 $S_n$ 进行分类讨论：

$$S_n = (q^0+q^1+q^2+\dotsb+q^{n-1}) \bmod p \mskip0.8mu (n \geqslant 1)$$

当 $n=1$ 时，有

$$S_1=q^0=1$$

当 $n>1$ 时，有

$$S_n=\begin{cases} \left[\left(1 + q^{\tfrac{n-1}{2}}\right) \cdot S_{\tfrac{n-1}{2}} + q^{n-1}\right] \bmod p &(n\text{为奇数})\\ \left(1 + q^{\tfrac{n}{2}}\right) \cdot S_{\tfrac n 2} \bmod p &(n\text{为偶数}) \end{cases}$$

该算法求和部分时间复杂度为 $O(\log n)$，若求解 $q^n$ 的算法时间复杂度为 $O(x)$，则总时间复杂度为 $O( x\log n)$。

代码实现如下：

// q: 公比; n: 项数; mod: 模; qpow: 求幂函数; qmul: 求积函数
long long sum(long long q, long long n) {
	if (n == 1) {
		return 1;
	}
	// 即 n % 2 == 1
	if (n & 1) {
		// 整型除法具有向下取整的特点
		// n 为奇数时, (n - 1) / 2 === n / 2, 可据此简化代码
		// n >> 1 即 n / 2
		return (qmul(1 + qpow(q, n >> 1), sum(q, n >> 1)) + qpow(q, n - 1)) % mod;
	} else {
		return qmul(1 + qpow(q, n >> 1), sum(q, n >> 1));
	}
}

2.2.2 求解 $q^{n}$ ——快速幂取模

解决了等比求和的问题，就只剩求解 $q^n$ 了。 此处选择了位运算的快速幂取模算法，时间复杂度为 $O(\log n)$。

代码实现如下：

// base: 底数; power: 指数; mod: 模; qmul: 求积函数
long long qpow(long long base, long long power) {
	// 对指数的底先取模，减少运算次数
	base %= mod;
	long long ans = 1;
	while (power) {
		// 即 power % 2 == 1
		if (power & 1) {
			ans = qmul(ans, base);
		}
		base = qmul(base, base);
		// 即 power /= 2
		power >>= 1;
	}
	return ans;
}

2.2.3 int64 的乘法模运算——位运算法

这部分优化直接借用 [参考文章4]，将大数乘法的时间复杂度从反复翻倍法（按位相乘）的 $O(\log n)$ 降到了 $\color{red}{O(1)}$。

代码实现如下：

// a, b: 乘数; mod: 模
long long qmul(long long a, long long b) {
	// 前两句处理了负数和数值过大的情况
	a = (a % mod + mod) % mod;
	b = (b % mod + mod) % mod;
	long long c = a * (long double) b / mod;
	long long ans = a * b - c * mod;
	// 若结果不在 [0, mod) 区间则调整一下
	if (ans < 0) {
		ans += mod;
	} else if (ans >= mod) {
		ans -= mod;
	}
	return ans;
}

3 实现

3.1 求和公式

完整代码实现如下：

总时间复杂度为 $O(\log n)$。

#include<stdio.h>
#define mod 100000007

long long qmul(long long a, long long b) {
	a = (a % mod + mod) % mod;
	b = (b % mod + mod) % mod;
	long long c = a * (long double) b / mod;
	long long ans = a * b - c * mod;
	if (ans < 0) {
		ans += mod;
	} else if (ans >= mod) {
		ans -= mod;
	}
	return ans;
}

long long qpow(long long base, long long power) {
	base %= mod;
	long long ans = 1;
	while (power) {
		if (power & 1) {
			ans = qmul(ans, base);
		}
		base = qmul(base, base);
		power >>= 1;
	}
	return ans;
}

int main() {
	int t;
	scanf("%d", &t);
	long long q, n;
	while (t--) {
		scanf("%lld %lld", &q, &n);
		if (q == 1) {
			printf("%lld\n", n % mod);
		} else {
			printf("%lld\n", (qpow(q, n) - 1) * qpow(q - 1, mod - 2) % mod);
		}
	}
	return 0;
}

提交结果如下：

7ms，低到惊人的耗时！

3.2 二分递归

完整代码实现如下：

总时间复杂度为 $O(\log^2 n)$

#include<stdio.h>
#define mod 100000007

long long qmul(long long a, long long b) {
	a = (a % mod + mod) % mod;
	b = (b % mod + mod) % mod;
	long long c = a * (long double) b / mod;
	long long ans = a * b - c * mod;
	if (ans < 0) {
		ans += mod;
	} else if (ans >= mod) {
		ans -= mod;
	}
	return ans;
}

long long qpow(long long base, long long power) {
	base %= mod;
	long long ans = 1;
	while (power) {
		if (power & 1) {
			ans = qmul(ans, base);
		}
		base = qmul(base, base);
		power >>= 1;
	}
	return ans;
}

long long sum(long long q, long long n) {
	if (n == 1) {
		return 1;
	}
	if (n & 1) {
		return (qmul(1 + qpow(q, n >> 1), sum(q, n >> 1)) + qpow(q, n - 1)) % mod;
	} else {
		return qmul(1 + qpow(q, n >> 1), sum(q, n >> 1));
	}
}

int main() {
	int t;
	scanf("%d", &t);
	long long q, n;
	while (t--) {
		scanf("%lld %lld", &q, &n);
		if (q == 1) {
			printf("%lld\n", n % mod);
		} else {
			printf("%lld\n", sum(q, n));
		}
	}
	return 0;
}

提交结果如下：

相较于等比求和通项公式法，二分法耗时有所增加，但依旧相当快速且适用范围更广。

4 参考文章

1. 《poj 1845 Sumdiv 数论--等比数列和(逆元或者递归)》
2. 《等比数列求和，二分递归》
3. 《费马小定理（易懂）》
4. 《小技巧1——长整型：64位整数的乘法模运算》