E – Divisible Substring 解説（AtCoder Beginner Contest 158）
問題概要
長さ N の数 S が与えられる。
素数 P で割り切れるような区間の選び方は何通りあるか？
制約
- \(1 \leq N \leq 2 \times 10^5\)
- \(2 \leq P \leq 10000\)
考え方
N が大きいので、普通に \(O(N^2)\) で区間を全探索し、P で割り切れるか確認するのは間に合いません。何らかの工夫が必要です。
P が比較的小さいことに注目
「ある数 P での mod を考える」ような問題は、P が小さければ
- 「a \(\equiv\) b (mod P)」\( \Leftrightarrow\)「a-b \(\equiv\) 0 (mod P)」
- P で割った余りが x になる数が何個あるか
を意識してみると解けることが多い印象があります。
特に前者を考えると「ある区間 sum が P で割り切れる」を次のように言い換えられることが分かります。
- 「sum \(\equiv\) a-b 」かつ 「a \(\equiv\) b (mod P) 」が成立する
つまり、「sum を差の形で表すことができ、その時の2つの剰余が等しければ、sum はPで割り切れる」ということです。
どうやって差の形で表すか
累積和的に右から見たときの値を考えてみましょう
例えば、N=4, P=3, S=3543 ならば、
- sum[4] = 0
- sum[3] = 3
- sum[2] = 43
- sum[1] = 543
- sum[0] = 3543
などとすると、これらを組み合わせることで、すべての区間を差の形で表現できます。
- 区間 [ i , j ) の値 = \(\frac{sum[i]-sum[j]}{10^{j-i-1}}\)
たとえば、区間 [0, 3) の値 354 は
- \(354 = \frac{3543 – 3}{10^2} = \frac{sum[0]-sum[3]}{10^2}\)
などと、\(10^{j-i-1}\) 分の1されてはいますが、差の形で表現できました。
割り算しても問題ないのか
合同式の性質
一般に P を法とした合同式の性質として、
- 和：「\(a \equiv b\) かつ \(c \equiv d\)」 \( \Rightarrow\) 「\( a+c \equiv b+d\)」
- 差：「\(a \equiv b\) かつ \(c \equiv d\)」 \( \Rightarrow\) 「\( a-c \equiv b-d\)」
- 積：「\(a \equiv b\) かつ \(c \equiv d\)」 \( \Rightarrow\) 「\( a×c \equiv b×d\)」
が成立します。
割り算の場合は特殊で、
- 商：「\(ab \equiv ac\) かつ \(a\) と P が互いに素」 \( \Rightarrow\) 「\( b \equiv c\)」
というのが成り立ちます。
\(10^k\) で割っても互いに素なら関係ない
以上より、積に関する法則から
- 「\(\frac{sum[i]-sum[j]}{10^{j-i-1}} \equiv 0\)」\( \Rightarrow\)「 \(sum[i]-sum[j] \equiv 0\)」
また、P と \(10^{j-i-1}\) が互いに素ならば、商に関する法則より
- 「 \(sum[i]-sum[j] \equiv 0\)」\( \Rightarrow\) 「\(\frac{sum[i]-sum[j]}{10^{j-i-1}} \equiv 0\)」
これら2つを合わせることで
- 「\(\frac{sum[i]-sum[j]}{10^{j-i-1}} \equiv 0\)」\( \Leftrightarrow\)「 \(sum[i]-sum[j] \equiv 0\)」
が成立することが分かりました。
つまり \(sum[i] \equiv sum[j] \) となるペアの数を求めれば良いことが分かります。
P と \(10^{j-i-1}\) が互いに素でないときは別に考える
P が 2 または 5 の時は、P と \(10^{j-i-1}\) が互いに素ではありません。
この時は実は、下一桁を見るだけで割り切れるか判定することができます。
例えば、24 は 一桁目の4 が 2の倍数なので 2 で割り切れます。また、35 は 一桁目の5 が 5 の倍数なので 5 で割り切れます
解き方
P が 2 or 5 の時
- 文字列 S の左から i 桁目 が P で割り切れた時、i+1 個 P で割り切れる区間がある(0-indexed)
という性質を利用します。これを i=0 から N-1 まで行えば良いので、計算量は \(O(N)\) です。
例えば、P=5, S=3653 の時、左から 2 桁目の 5 は P で割り切れて、
- 5
- 65
- 365
の 2+1 個も全て P で割り切れます。
P が 2 でも 5 でも無い時
P と \(10^{j-i-1}\) が互いに素になります。
- sum[i] := S の [i, N) の区間が表す数
というようにすると、
- 「\(\frac{sum[i]-sum[j]}{10^{j-i-1}} \equiv 0\)」\( \Leftrightarrow\)「 \(sum[i]-sum[j] \equiv 0\)」
となるので、\(sum[i] \equiv sum[j] \) となるペアの数を求めれば良いです。
この数は、例えば x mod P となる数を m[x] などとすると、
- m[x] *(m[x]-1) /2
で求めることができます。
C++ での実装例
sum[i] の計算のために、\(10^{N-i}\)%P を表す ten[i] を前計算しています。
#include <bits/stdc++.h>
#define ALL(obj) begin(obj), end(obj)
using namespace std;
using ll = long long;
using ull = unsigned long long;
int main() {
ll N, P;
string S;
cin >> N >> P;
cin >> S;
// P が 2 or 5 の時
if (P == 2 || P == 5) {
ll ans = 0;
for (int i = 0; i < N; i++) {
if ((int)(S[i] - '0') % P == 0) {
ans += i + 1;
}
}
cout << ans << endl;
return 0;
}
// P が 2 でも 5 でも無い時
vector<ll> sum(N + 1, 0), ten(N + 1, 1);
for (int i = N - 1; i >= 0; i--) { // ten の前計算
ten[i] = ten[i + 1] * 10 % P;
}
for (int i = N - 1; i >= 0; i--) { // sum の計算
sum[i] = ((ll)(S[i] - '0') * ten[i + 1] % P + sum[i + 1]) % P;
}
map<ll, ll> m; // 剰余が同じものを数える
for (auto i : sum) {
m[i]++;
}
ll ans = 0;
for (auto i : m) {
ans += i.second * (i.second - 1) / 2;
}
cout << ans << endl;
return 0;
}