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);
}