AtCoder ABC304-F Shift Table を 2 通りの方法で解く

# はじめに

[東京海上日動プログラミングコンテスト2023（AtCoder Beginner Contest 304） \- AtCoder](https://atcoder.jp/contests/abc304) 2023/06/03 お疲れさまでした。


[F \- Shift Table](https://atcoder.jp/contests/abc304/tasks/abc304_f) シフト表を考える問題が面白く感じました。コンテスト後に 2通りの方法で解きました。面白くて図示してみました。その流れで記事にまとめてみた、というものです。


# TL;DR

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/118172/a2167597-1304-135b-82d2-f4e5c9c3a71d.png)

問題の紹介は省略します。公式サイトをどうぞ。

# 解き方

入力例3 をもとに考えます。12日の高橋君のシフト表が分かっています。

```
12
####.####.##
```
出力例
```
19
```

## 周期と、周期ごとの候補数

12日の約数は 1, 2, 3, 4, 6, 12 です。このうち最後の 12 は繰り返しなしですので条件を満たしません。残り 5つが周期になります。

そして周期ごとの候補数は簡単に求まります。たとえば 4日周期を考えます。高橋君の出勤日を 4日ごとに折り返すと次のようになります。

```
####
.###
#.##
```

* 高橋君は 4日周期の 1, 2 日めにお休みすることがあります。青木君はこの日に出勤しないといけません。場合の数は増えません。
* 高橋君は 4日周期の 3, 4 日めは必ず出勤します。青木君はこの日は出勤・欠勤どちらでも良いです。 2通りどちらでも良いので、場合の数は 2倍ずつ増えていきます。

青木君シフトの候補は `##??` 、計 2 * 2 = 4通りとなります。

Rust だと `i` が i 日周期を表すときに、それぞれの場合の数はこんな感じで求まります。

```rust
let mut x = 1;
let j = n / i;
for i0 in 0..i {
    if (0..j).all(|j0| s[i * j0 + i0] == '#') {
        x = (x * 2) % MOD;
    }
}
```

|周期|シフト|候補数　|
|---|---|---|
|1|# * 12|$2^0=1$|
|2|## * 6|$2^0=1$|
|3|##? * 4|$2^1=2$|
|4|##?? * 3|$2^2=4$|
|6|???##? * 2|$2^4=16$|

答えが近づいた気がします。でも正解 19 につなげるには材料が足りなさそうです。


## 重複を取り除きたい

この表で示した集合は、ベン図のように重なり合っています。残念ながら単純に足し算できません。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/118172/8cf64700-586a-13df-ce4c-7065716e95ec.png)

重複の数が少なければ `set` に答えを入れれば済むでしょう。でも、「998244353 で割った余りを求めてください」のように多くの組み合わせを扱うこの問題に対しては力不足です。

周期の短い方から埋めるか、長い方から埋めるかの 2択になると思います。


## 方針1: 周期の短い方から埋める

表をもう一度、周期の小さな方から眺めていきます。

|周期|シフト|候補数|
|---|---|---|
|1|# * 12|1|
|2|## * 6|1|
|3|##? * 4|2|

* 周期 1日の場合は `# # # # # # # # # # # #` 1通りです。
* 周期 2日の場合は `## ## ## ## ## ##` です。これは周期 1日で現れたものと同じです。2日にしたことで追加できるシフトはありません。
* 周期 3日の場合は `### ### ### ###` と `##. ##. ##. ##.` です。前者は周期 1日で現れたものと同じです。3日にしたことで追加できるシフトは後者 1通りです。

このように「以前の周期でも表現できるもの」 = 「約数となる周期のもの」の個数を除くと、重複を除くことができそうです。

続けて表を完成させてみます。

|周期|シフト|候補数　|重複 (約数)|この周期で追加できる候補数|
|---|---|---|---|---|
|1|# * 12|$2^0=1$||$2^0=1$|
|2|## * 6|$2^0=1$|周期1|$2^0-1=0$|
|3|##? * 4|$2^1=2$|周期1|$2^1-1=1$|
|4|##?? * 3|$2^2=4$|周期1, 2|$2^2-(1+0)=3$|
|6|???##? * 2|$2^4=16$|周期1, 2, 3|$2^4-(1+0+1)=14$|
|計||||$1+0+1+3+14=19$|

19 になりました。


### 方針1 で実装

この表を `let mut v = vec![(1, 1), (2, 0), (3, 1), (4, 3), (6, 14)];` のように周期と追加候補数のペアを持つ形で組み立ててみました。実装例です。

```rust
use proconio::{input, marker::Chars};

fn main() {
    input! {
        n: usize,
        s: Chars,
    }

    const MOD: usize = 998244353;
    let mut v = vec![];
    for i in 1..=(n / 2) {
        if n % i > 0 {
            continue;
        }

        let mut x = 1;
        let j = n / i;
        for i0 in 0..i {
            if (0..j).all(|j0| s[i * j0 + i0] == '#') {
                x = (x * 2) % MOD;
            }
        }
        let x0 = v
            .iter()
            .filter(|&(i0, _)| i % i0 == 0) // 登録済みの約数を集める
            .map(|&(_, x)| x)
            .sum::<usize>()
            % MOD;
        x = (x + MOD - x0) % MOD;
        v.push((i, x));
    }

    let result = v.iter().map(|&(_, x)| x).sum::<usize>() % MOD;
    println!("{}", result);
}
```

約数列挙用のライブラリーを持っている方も多いと思います。今回は最大でも $10^5$ だからと、べたっとその場で計算しました。

公式解説と似た方針のつもりです。公式解説では modint を使い、事前に 2の累乗を用意するなど、よりコンパクトに書いています。


## 方針2: 周期の長い方から絞り込む

もういちどもとのベン図を見てみます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/118172/8cf64700-586a-13df-ce4c-7065716e95ec.png)

長い周期 6, 4 を調べて、重複部分の周期 2 を消せば、16 + 4 - 1 = 19 と答えが見つかりそうです。周期 3, 1 は調べなくても良さそうです。

長い周期は、約数のうち素数 1つだけ割ったものとなります。6 = 12 / 2, 4 = 12 / 3 と、それぞれ 2, 3 で割ったものです。2つ以上の素数で割ったものは、ほかの集合に含まれます。

重複部分は長い周期の最大公約数で求まります。2つの素数から組み立てられる整数なら、これで求まります。

では、3つ以上の素数から組み立てる場合はどうか、というと「包除原理」によって同じように扱えます。全部が重なるところまで繰り返し、プラスとマイナスを行います。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/118172/9688b6ac-3f2e-76da-2723-e806becd1a89.png)

3つの素数 '2, 3, 5' から 'k' 個選ぶ方法は組み合わせ $_3C_k$ です。配列の combination をライブラリを使って順に調べ、それぞれの積で割り算 $60/Π$ すれば、対象全ての周期長さが得られます。

### 方針2 で実装

素数かどうか調べるのに、方針 1 よりひと手間かかります。競技プログラミングしているとよく現れる処理ですから、きっとコピペくらいの手間でしょう。
その後の実装量は方針 1 とあまり変わりません。

```rust
use itertools::Itertools;
use proconio::{input, marker::Chars};

fn main() {
    input! {
        n: usize,
        s: Chars,
    }

    const MOD: usize = 998244353;

    let mut primes = vec![];
    {
        let mut n = n;
        let mut k = 2;
        while k * k <= n {
            if n % k == 0 {
                primes.push(k);
                while n % k == 0 {
                    n /= k;
                }
            }
            k += if k == 2 { 1 } else { 2 };
        }
        if n > 1 {
            primes.push(n);
        }
    }

    let mut result = 0usize;
    for k in 0..primes.len() {
        for v in primes.iter().combinations(k + 1) {
            let i = v.iter().fold(n, |acc, &&x| acc / x);
            let mut x = 1;
            let j = n / i;
            for i0 in 0..i {
                if (0..j).all(|j0| s[i * j0 + i0] == '#') {
                    x = (x * 2) % MOD;
                }
            }

            if k % 2 == 0 {
                result = (result + x) % MOD;
            } else {
                result = (result + MOD - x) % MOD;
            }
        }
    }

    println!("{}", result);
}
```

# 最後に

本番中に解けるようになりたいです……。