読者です 読者をやめる 読者になる 読者になる

Lilliput Steps

小さな一歩から着実に. 数学やプログラミングのことを書きます.

Codeforces 359C - Prime Number

問題文

Prime Number

概要

$x$ を素数とする.

$\dfrac{1}{x^{a_1}} + \dfrac{1}{x^{a_2}} + \cdots + \dfrac{1}{x^{a_n}}$ を$\dfrac{s}{t}$ とするとき(ここで$t=x^{a_1 + a_2 + \cdots + a_n}$), $\text{gcd}(s, t)$ mod $10^9+7$ を求めよ.

$1 \leqq n \leqq 10^5,\ 1 \leqq a_i \leqq 10^9,\ 2 \leqq x \leqq 10^9$


解法

明らかに答えは$x$ のべきとなる. 具体的には, $S=a_1+a_2+\cdots+a_n$ とすると $x^{ \min \{ S - a_i \} }$ が答えとなりそうだが, $S - a_i$ が$x$ の倍数個あるとするとそれらは$x$ 個まとめることで$S - a_i + 1$ 個にすることができる.

このようにまとめていく作業はたかだか$O(n)$ でできるが, ある$S - a_i$ の個数を求めるのには$O(n \log n)$ 時間の処理(ソートとかstd::map とか) が必要なので全体の計算量は$O(n \log n)$ である.

コード

#include <cstdio>
#include <vector>
#include <algorithm>

#define MOD (1000000007)

using namespace std;

typedef long long lint;

lint pow(lint a, lint n)
{
    lint ret = 1, b = a;
    
    while (n){
        if (n & 1) ret = (ret * b) % MOD;
        b = (b * b) % MOD;
        n >>= 1;
    }
    return (ret);
}

int main()
{
	int n;
	int x;
	lint sum = 0;
	vector<lint> a;
	
	scanf("%d %d", &n, &x);
	
	for (int i = 0; i < n; i++){
		int y;
		scanf("%d", &y);
		a.push_back(y);
		sum += y;
	}
	for (int i = 0; i < n; i++) a[i] = sum - a[i];
	
	sort(a.rbegin(), a.rend());
	
	while (1){
		lint target = a[a.size() - 1];
		int cnt = 0;
		while (a.size() && a[a.size() - 1] == target){
			++cnt; a.pop_back();
		}
		
		if (cnt % x){
			printf("%lld\n", pow(x, target));
			break;
		}
		else {
			for (int i = 0; i < cnt / x; i++) a.push_back(target + 1);
		}
	}
	
	return (0);
}