きままにものづくり

日々の気付いたことなんかを書いてます。

Codeforces Round #191 (Div. 2) C MagicFive

問題

n桁の整数を持つ集合sがある。集合sは整数の集合a(1<=a<=10^5)をk回(1<=k<=10^9)つなぐことで、作られる。この中で順番は変更せずに、任意の桁を削除することで5で割れる数にしなさい。先頭にゼロを含んでいてもよい。このような集合を作る、削除の組み合わせの総数を求めなさい。
答えは、1000000007のmodをとって出力しなさい。

解説

この問題は、等比数列の公式とモジュロ逆数を用いることで解くことができます。
まずは、具体例からみていきます。

Input
1256
1
Output
4

この場合、集合sは1256になります。
全ての桁に対して、消すか消さないかで全探索すると計算量はO(2^n)になります。この例では、n=5なので短時間で求めることができます。しかし、nはa*kなので、nは最大で10^{14}になります。これでは、とても間に合いません。
他の方法を考えてみましょう。
5または0以外が一桁目にあると、5で割ることはできません。なので、この例ですと6は常に消さなければなりません。この条件は言い換えると、5または0が一桁目になければ、5で割れないと言い換えることができます。
これで、選択肢として残るのは1を消すか、2を消すか、どちらも消さないか、はたまたどちらも消すかの4択となります。
消す状態を0、消さない状態を1とすると、下のような表になります。

1 2
0 0
0 1
1 0
1 1

これは、2ビットの真理値表と等価になります。つまり、ビット数をbとすると2^bの組み合わせがあることになります。
よって、このサンプルの答えは4になります。

5または0がある位置から左の要素の数をlとすると、その時の組み合わせは2^lになります。lが0の時は、その要素を残すという選択肢しかないため2^0、つまり1になります。lは、集合sを左から、つまりindex0から線形探索することで求めることができます。この例だと、5または0はひとつしか含まれいません。複数含まれている場合を考えてみましょう。

Input
152063
1
Output
10

この例ですと、5を一桁目にする時は、1を消すかどうかしか選択できないので2^1になります。一桁目には、もうひとつ0を選ぶことができます。この時の選択肢は1,5,2と3つあるので2^3となります。5を一桁目にする場合と0を一桁目にする場合とで、重なりあう選択肢はひとつもないので、足し合わせることができます。
よって、答えは2^1+2^3 = 10となります。

sを線形探索し、0または5があったらその位置に対応した組み合わせを足し合わせていくという方法で答えにたどり着けそうです。このアルゴリズムの計算量はO(n)になります。全探索の場合と比較するといくらかましになりました。しかし、nは最大で10^14なのでこのアルゴリズムでは不十分です。

sで線形探索すると、要素数が大きくなってしまいます。sはaの繰り返しになっているので、aを用いることでsの全ての組み合わせを求めることができそうです。
二個目の例のkが2以上の例を考えてみましょう。

Input
152063
3
Output
41610

sは152062152063となります。これの答えは、線形探索を用いたアルゴリズムを用いると2^1+2^3+2^7+2^9+2^{13}+2^{15} = 41610となります。ここで、式をよくみるとaの時の答えに、2^6をかけているということに気づきます。式変形すると、(2^1+2^3)*(2^0+2^6+2^{12})になります。aの組み合わせの数をS、aの要素数をLとして、一般的な形で表すと全ての組み合わせ数Cは以下の式になります。ただし、つなぎあわせる回数kは以下の式ではKとしています。
<br />
\begin{eqnarray}<br />
C = S * \sum_{k=0}^{K-1}2^{Lk}<br />
\end{eqnarray}<br />
シグマの式は、等比数列となっているので
<br />
\begin{eqnarray}<br />
C = S*\frac{2^{Lk}-1}{2^L-1}<br />
\end{eqnarray}<br />
となります。
これなら、aの組み合わせ数を求めれば式に代入するだけなので、計算量はO(a)になります。aは最大で10^5なので、これなら時間内に解くことができます。
使用する言語によってはこれで答えを求めることができます。しかし、オーバーフローしてしまう言語では工夫が必要になります。理由は、modの計算は除算に対しては閉じていないからです。
例を見てみましょう。
 8 mod 5 = 3
 4 mod 5 = 4
 \frac{8}{4} mod 5 = 2
modの計算を先に計算すると\frac{3}{4}となってしまい、答えがかわってしまいます。
modでの除算を実現するためには、モジュラ逆数を使います。これは、拡張ユークリッドの互除法オイラーの定理を用いることで求めることができます。

コード

#include <vector>
#include <string>
#include <iostream>
#include <fstream>
#include <sstream>
#include <list>
#include <algorithm>
#include <sstream>
#include <set>
#include <cmath>
#include <map>
#include <stack>
#include <queue>

#define MAX_V (10000001)
#define MAX_N (100001)
#define MOD 1000000007

using namespace std;

long long
modPow(long long x,long long n,long long m)
{
	if (n==0) return 1;
	long long res = modPow((x*x)%m, n/2, m);
	if (n&1) res = (res*x)%m;
	return res;
}

long long
extgcd(long long a,long long b,long long &x,long long &y)
{
	long long d = a;
	if (b!=0) {
		d = extgcd(b,a%b,y,x);
		y -= (a/b)*x;
	}
	else {
		x = 1;
		y = 0;
	}
	return d;
}

long long
inv(long long x)
{
	return modPow(x, MOD-2, MOD);
}

long long
solve(string str,int k)
{
	long long single_ans=0;
	for (int i=0; i<str.size(); ++i) {
		if (str[i]=='0'||str[i]=='5') {
			single_ans = (single_ans + modPow(2, i, MOD))%MOD;
		}
	}
	long long q = modPow(2, str.size(), MOD);
	long long sum = ((modPow(q, k, MOD)-1)*inv(q-1))%MOD;
	return single_ans*sum%MOD;
}

int main(void) {
	string str;
	int n;

	cin >> str;
	cin >> n;

	long long ans = solve(str,n);
	cout << ans;
	return 0;
}