Pythonのジェネレータを使って競プロを遊ぶ


# この記事は何
Pythonで競プロを遊ぶときにジェネレータをたまに使っているのですが、使用感が良いので紹介してみようという記事です。個人的に感じたメリットとしては、
* メイン部分の実装が簡潔になる
* Python特有の機能を使えている気分になれてうれしい（？）

などがあります。
Pythonでの競プロ実装する際の一方針を紹介するものであり、計算速度が改善したりするわけではないです。

動作はPython3.8.6で確認しています。
不備や不正確な点などがありましたら、ご指摘いただけるとうれしいです。

# ジェネレータとは？
pythonのジェネレータに関してざっくり説明します。
詳しい説明は[Documentation](https://docs.python.org/ja/3/howto/functional.html#generators
)や他の方の記事を参照いただければと思います。

関数の`return`を`yield`に置き換えるとジェネレータになります。
関数は値を`return`した時点で処理を終了します。一方、ジェネレータでは値を`yield`した時点の状態が保持され、再びジェネレータが呼び出されるとその続きから処理が再開されます。
関数が**値**を返すのに対して、ジェネレータは一連のデータをひとつずつ返すオブジェクトである**イテレータ**を返してくれます。


全競プロpythonistaが使っている`map()`もジェネレータで同じ機能が実装できます。
```Python
def mymap(func, value):
    for v in value:
        yield func(v)


inp = "1 1 2 3 5 8"
A = list(mymap(int, inp.split()))
B = list(map(int, inp.split()))
print(A)    # -> [1, 1, 2, 3, 5, 8]
print(B)    # -> [1, 1, 2, 3, 5, 8]
```
返したい一連の値をジェネレータで具体的に設定できる、という感じです。

ジェネレータで定義した一連のデータは、`list()`でリスト化したり、`max()`で最大値を求めたり、後述するようにfor文で値を逐次取り出すことができます。値をひとつずつ取得したいときは`next()`関数を使用します。


# 競プロにおける使用例
簡単な例4つほど

## 1. 条件分岐（if文）をいい感じにまとめる
[ABC135D Digits Parade](https://atcoder.jp/contests/abc135/tasks/abc135_d)を例に紹介します。（軽いネタバレがあります）

 
　
　
　
　


桁の一部が`?`で隠された整数$S$が与えられるので、隠された部分を数字に置き換えたときに条件を満たすものを数え上げる問題です。　
この問題では$S_i$が`?`のときは`0 ~ 9`の数字に置き換えて、$S_i$が数字のときはそのまま用いてdpテーブルを更新していくことで解くことができます。素直に実装するとこんな感じになると思います。
```Python
for i in range(len(S)):
    for j in range(13):
        if S[i] == "?":
            for k in range(10):
                dp[i+1][ ] += dp[i][k]    # dp遷移
        else:
            dp[i+1][ ] += dp[i][int(S[i])]    # dp遷移
```
$S_i$の値に応じて処理が少し変わるため、同じようなdp遷移を2回書くことになっていますね。
ジェネレータを用いてif文の部分をまとめると、下のようになります。
```Python
def gen(s):
    if s == "?":
        yield from range(10)
    else:
        yield int(s)


for i in range(len(S)):
    for j in range(13):
        for k in gen(S[i]):
            dp[i+1][ ] += dp[i][k]    # dp遷移
```
ジェネレータ`gen`は、引数sが`?`のときは整数`0 ~ 9`をひとつずつ返してくれ、sが数字の場合はsを整数に変換したものを返してくれます。メインのforループの部分を少しすっきりさせることができましたね（[実装例](https://atcoder.jp/contests/abc135/submissions/32500893)）。
`yield from [iterableなオブジェクト]`と書くことで、オブジェクトから要素をひとつずつyieldすることができます。

## 2. 部分集合の列挙
bitDPでたまに見る、$N$要素からなる集合の部分集合$S$の部分集合$T$を$O(3^N)$で列挙するやつです（問題例：[EDPC](https://atcoder.jp/contests/dp/tasks/dp_u)、[典型90](https://atcoder.jp/contests/typical90/tasks/typical90_as)）。

cppだと以下のように書けます。シンプルでいいですね
```cpp:cpp
for (int S=0; S<(1 << N); ++S) {
	for (int T=S; T>0; T=(T-1)&S) {
 		dp[S] = ...      // dp遷移
    }
}
```
pythonの`range()`では、上コード2行目のような条件式や増減式をワンライナーで実装することはできません。while文を用いて以下のように書くことになると思います。
```Python
for S in range(1 << N):
    T = S
    while T:
        dp[S] = ...      # dp遷移
        T = (T - 1) & S
```
これでは頭が爆発してしまいます（断言）。
ジェネレータで条件式などを分けて実装すると、下のようになります。
```Python
def subset_of_(S):
    T = S
    while S:
        yield S
        S = (S - 1) & T


for S in range(1 << N):
    for T in subset_of_(S):
        dp[S] = ...      # dp遷移
```
メインのforループがシンプルになっておりdp遷移に集中して実装できそうな気がします。
ジェネレータ`subset_of_`はfor文を用いることで以下のように$S$の部分集合$T$を降順に列挙します。
```Python
def subset_of_(S):
    T = S
    while S:
        yield S
        S = (S - 1) & T


S = int("1101", 2)
for T in subset_of_(S):
    print(bin(T)[2:].zfill(4), T)
 
"""
出力
1101 13
1100 12
1001 9
1000 8
0101 5
0100 4
0001 1
"""
```
ビット演算系はジェネレータを用いてライブラリ化しやすいと思います。


## 3. 周期的に値を生成する
周期3の場合を考えてみます（下の実装例を見たほうが早いかもです）。
$3k$番目に`a`、$3k+1$番目に`b`、$3k+2$番目に`c`を取り出したい状況を考えます（$k$は非負整数です）。[この問題](https://atcoder.jp/contests/abc238/tasks/abc238_g)で使いたくなりました（[提出](https://atcoder.jp/contests/abc238/submissions/32508132)）。

```Python
def gen():
    a = 1
    b = 4
    c = 9
    while True:
        yield a
        yield b
        yield c


g = gen()
print(next(g))      # -> 1
print(next(g))      # -> 4
print(next(g))      # -> 9
print(next(g))      # -> 1

for i, x in enumerate(g):
    print(x, end=", ")   # ->  4, 9, 1, 4, 9, 1, ...
    if i >= 10:
        break
```
`g = gen()`で生成したイテレータ`g`に対して、`next()`を適用することで`a`、`b`、`c`、`a`、`b`、`c`、...と繰り返し生成してくれます。yieldに到達した段階の状態が保持されるため、このような機能が実現できるわけですね。（上の例だとfor文を使うと無限ループになるので注意です。）


## 4. [Xorshift](https://ja.wikipedia.org/wiki/Xorshift)
疑似乱数列を生成するアルゴリズムです。

```Python
def xor64(seed=88172645463325252):
    x = seed
    mask = (1 << 64) - 1
    while 1:
        x = x ^ (x << 13) & mask
        x = x ^ (x >> 7)
        x = x ^ (x << 17) & mask
        yield x


rnd = xor64()
for _ in range(5):
    print(next(rnd))
"""
出力
8748534153485358512
3040900993826735515
3453997556048239312
16431732851926010853
8204724074003728306
"""
```

関数を用いてxorshiftを実装しようとすると上のコードでいう`x`をグローバル変数にしたりする必要がありそうですが、ジェネレータの実装では不要になります。

pythonの整数はオーバーフローしないため、毎回bitmaskしているのがあまりスマートではないですね。。パフォーマンスもそこまで良くないので乱数を扱いたいならpythonのrandomモジュールを使えばいいとは思います。


余談ですが、どうしてもxorshiftを自前実装したい場合、for文で使わないのであればジェネレータよりもクロージャを使うほうがちょっとスマートかもしれません。乱数を生成するときに`next()`を毎回書く必要がないので。


<details><summary>クロージャを用いた実装例</summary><div>

```Python
def xor64(seed=88172645463325252):
    x = seed
    mask = (1 << 64) - 1

    def inner():
        nonlocal x
        x = x ^ (x << 13) & mask
        x = x ^ (x >> 7)
        x = x ^ (x << 17) & mask
        return x
    return inner


rnd = xor64()
for _ in range(5):
    print(rnd())
"""
出力
8748534153485358512
3040900993826735515
3453997556048239312
16431732851926010853
8204724074003728306
"""
```

</div></details>


# おわり

無限ループ、気をつけようね

# 参考
https://qiita.com/tomotaka_ito/items/35f3eb108f587022fa09

https://qiita.com/keitakurita/items/5a31b902db6adfa45a70

https://docs.python.org/ja/3/howto/functional.html#generators

https://ja.wikipedia.org/wiki/Xorshift