LoginSignup
17
29

More than 3 years have passed since last update.

@Kiri8128(Kiri8128)

可変長ループを書こう！

Last updated at Posted at 2021-03-05

はじめに

突然ですが、あなたは多重ループは書けますか？

？？？：「書けますー！」

test.py
cnt = 0
for i0 in range(3):
    for i1 in range(3):
        for i2 in range(3):
            for i3 in range(3):
                ##### ここに処理を書く #####
                cnt += 1
                #####
print(cnt) # 81

これは $4$ 重ループになってますが、これの可変なやつ（詳細は後述）をうまく処理しようというのがこの記事の趣旨です。言語は Python 、特に PyPy での実行を想定しています。

ループを書くことが目的なので「処理」のところは何でもよいのですが、ここではループした回数をカウントしています。実際の問題では、 ans += calc(i0, i1, i2, i3) のような数え上げを計算したり、 ans = max(ans, calc(i0, i1, i2, i3)) のような更新をしたり、あるいは DP テーブルを更新したりするイメージです。

可変長ループとは

だいたいこんなのができるものを可変長ループと呼んでいます 1。

  • 可変個（$K$ 個とします）の変数 $i_0,\ \cdots, i_{K-1}$ でループをする
  • $i_k$ は $[l_k,\ r_k)$ の範囲の整数をループする。 $l_k,\ r_k$ は $i_0,\ \cdots,\ i_{k-1}$ に依存してもよい
  • 変数間に満たすべき条件 $P(i_{k_0},\ \cdots,\ i_{k_{s-1}}) $ がいくつかある
  • 非再帰かつ 空間 計算量 $O(K)$
  • $i_0,\ \cdots,\ i_{k-1}$ が満たすべき条件を満足していない場合は $i_k$ 以降のループは回さない 2

いくつか例を挙げてみます。まずは愚直ループで書いてみるので、可変長にするにはどうすれば良いか考えながら見てみてください。

例1（愚直）

変数 $i_0,\ i_1,\ i_2,\ i_3$ が $0\le i_k < 2k+1$ を満たす整数を動くループを書いてください。
愚直に書くとこんな感じ。

test.py
cnt = 0
for i0 in range(1):
    for i1 in range(3):
        for i2 in range(5):
            for i3 in range(7):
                ##### ここに処理を書く #####
                cnt += 1
                #####

print(cnt) # 105

これぐらいなら Python の itertools を使うなどいろんな書き方があると思います3。

例2（愚直）

さっきのに加えて、「隣り合う変数の合計が $3$ の倍数になるものは除く」という条件を入れたらどうなるでしょうか。

test.py
cnt = 0
for i0 in range(1):
    for i1 in range(3):
        if (i0 + i1) % 3 == 0: continue
        for i2 in range(5):
            if (i1 + i2) % 3 == 0: continue
            for i3 in range(7):
                if (i2 + i3) % 3 == 0: continue
                ##### ここに処理を書く #####
                cnt += 1
                #####

print(cnt) # 31

これも itertools などでできそうですが、普通にループして全てのループ変数の組に対して条件を判定するやり方だと無駄なループができて計算量が大きくなってしまう可能性があります。具体的には、上の例だと (i0 + i1) % 3 == 0 の時点でその後の i2 や i3 を見る必要はないにも関わらず、ループを回してしまうことになりそうです。
なので上の continue のように、条件を満たさなくなった時点でそれ以降の変数のループを回さないようにしたいです。

例3（愚直）

次はループ範囲がそれより前の変数に依存する場合を考えてみましょう。
例えば $i_k$ のループ範囲が $[0,\ i_0+\cdots +i_{k-1}+2)$ のように表される場合です。この例では右側だけ動かしていますが、左側が前の変数に依存することもありえます。

test.py
cnt = 0
for i0 in range(1):
    for i1 in range(i0 + 2):
        if (i0 + i1) % 3 == 0: continue
        for i2 in range(i0 + i1 + 2):
            if (i1 + i2) % 3 == 0: continue
            for i3 in range(i0 + i1 + i2 + 2):
                if (i2 + i3) % 3 == 0: continue
                ##### ここに処理を書く #####
                cnt += 1
                #####
print(cnt) # 5

これのループ変数の個数を増やすと（例えば $K=10$ など）結構大変そうですね。
まあ書けなくはないんですけど。

test.py
cnt = 0
for i0 in range(1):
    for i1 in range(i0 + 2):
        if (i0 + i1) % 3 == 0: continue
        for i2 in range(i0 + i1 + 2):
            if (i1 + i2) % 3 == 0: continue
            for i3 in range(i0 + i1 + i2 + 2):
                if (i2 + i3) % 3 == 0: continue
                for i4 in range(i0 + i1 + i2 + i3 + 2):
                    if (i3 + i4) % 3 == 0: continue
                    for i5 in range(i0 + i1 + i2 + i3 + i4 + 2):
                        if (i4 + i5) % 3 == 0: continue
                        for i6 in range(i0 + i1 + i2 + i3 + i4 + i5 + 2):
                            if (i5 + i6) % 3 == 0: continue
                            for i7 in range(i0 + i1 + i2 + i3 + i4 + i5 + i6 + 2):
                                if (i6 + i7) % 3 == 0: continue
                                for i8 in range(i0 + i1 + i2 + i3 + i4 + i5 + i6 + i7 + 2):
                                    if (i7 + i8) % 3 == 0: continue
                                    for i9 in range(i0 + i1 + i2 + i3 + i4 + i5 + i6 + i7 + i8 + 2):
                                        if (i8 + i9) % 3 == 0: continue
                                        ##### ここに処理を書く #####
                                        cnt += 1
                                        #####
print(cnt) # 5964602

これをきれいに書くのを目標にしましょう 4。

他の方法

本記事で示す方法以外にもいくつか方法があります。これらが使えることもあるでしょう。

  1. 愚直にループを書く
    上の例で書いたようなやつです。ただこれだとループの深さなどが可変の問題には対応しづらいですね。

  2. itertools
    不要なところもループしてしまうので場合によっては計算量が増えてしまう可能性があります。

  3. 再帰 DFS
    再帰は言語によっては遅くなるので使いたくなかったりします。

  4. BFS
    メモリが $O(K)$ では収まらないです。

  5. 本記事の方法
    可変長対応、不要なループをしない 5 、非再帰、メモリは $O(K)$ 、のすべてを満足する方法です。

    アイデア

アイデアというほどでもないですが、本記事の実装方法です。

  • 各ループ変数に対して、「初期値」「終了値」を持つ。具体的には $i$ 番目のループ変数の初期値・終了値をそれぞれ L[i] および R[i] で表す
  • while 文をでループする

これだけです。

実装

さっきの例1

さっきの例で実装してみます。
chk 関数で満たすべき条件を設定しています。この例では特に NG 条件がないので必ず $0$ を返すようにしています。

test.py
def chk(i):
    if not i: return 0
    # i 番目の変数の条件（OK なら 0 を、NG なら 1 を返す）
    # i-1 番目までの変数に依存しても良い
    return 0

K = 4 # ループ変数の個数
Z = [0] * K
L = [0] * K
R = [0] * K
L[0], R[0] = 0, 1 # 半開区間 [L, R) をループする
i = 0
cnt = 0
while i >= 0:
    while i < K - 1:
        i += 1

        ##### 初期値・終了値の設定 #####
        L[i] = 0
        R[i] = 2 * i + 1
        #####

        Z[i] = L[i]
        if chk(i):
            break
    else:
        ##### ここに処理を書く #####
        cnt += 1
        #####

    Z[i] += 1
    while Z[i] >= R[i] or chk(i):
        if Z[i] < R[i]: Z[i] += 1
        while Z[i] >= R[i]:
            i -= 1
            if i < 0: break
            Z[i] += 1
        if i < 0: break

print(cnt) # 105

さっきの例2

ループ変数間に条件がある場合です。
chk 関数に条件を入れています。 $i$ 番目のループ変数の条件のところでは $i$ 番目までの変数に依存して決めて良いです。言い換えると、いくつかのループ変数に関係する条件については、その中で最も大きな番号のループ変数のところで判定をして、 NG ならすぐ次に移るようにしています。

test.py
def chk(i):
    if not i: return 0
    # i 番目の変数の条件（OK なら 0 を、NG なら 1 を返す）
    # i-1 番目までの変数に依存しても良い
    return 1 if (Z[i-1] + Z[i]) % 3 == 0 else 0

K = 4 # ループ変数の個数
Z = [0] * K
L = [0] * K
R = [0] * K
L[0], R[0] = 0, 1 # 半開区間 [L, R) をループする
i = 0
cnt = 0
while i >= 0:
    while i < K - 1:
        i += 1

        ##### 初期値・終了値の設定 #####
        L[i] = 0
        R[i] = 2 * i + 1
        #####

        Z[i] = L[i]
        if chk(i):
            break
    else:
        ##### ここに処理を書く #####
        cnt += 1
        #####

    Z[i] += 1
    while Z[i] >= R[i] or chk(i):
        if Z[i] < R[i]: Z[i] += 1
        while Z[i] >= R[i]:
            i -= 1
            if i < 0: break
            Z[i] += 1
        if i < 0: break

print(cnt) # 31

さっきの例3

$i$ 番目のループ変数の範囲が $i-1$ 番目までのループ変数に依存する場合です。
R[i] = sum(Z[:i]) + 2 のところで設定しています 6 。

test.py
def chk(i):
    if not i: return 0
    # i 番目の変数の条件（OK なら 0 を、NG なら 1 を返す）
    # i-1 番目までの変数に依存しても良い
    return 1 if (Z[i-1] + Z[i]) % 3 == 0 else 0

K = 4 # ループ変数の個数
Z = [0] * K
L = [0] * K
R = [0] * K
L[0], R[0] = 0, 1 # 半開区間 [L, R) をループする
i = 0
cnt = 0
while i >= 0:
    while i < K - 1:
        i += 1

        ##### 初期値・終了値の設定 #####
        L[i] = 0
        R[i] = sum(Z[:i]) + 2
        #####

        Z[i] = L[i]
        if chk(i):
            break
    else:
        ##### ここに処理を書く #####
        cnt += 1
        #####

    Z[i] += 1
    while Z[i] >= R[i] or chk(i):
        if Z[i] < R[i]: Z[i] += 1
        while Z[i] >= R[i]:
            i -= 1
            if i < 0: break
            Z[i] += 1
        if i < 0: break

print(cnt) # 5

さっきみたいに $K=10$ にしてみましょう。

test.py
def chk(i):
    if not i: return 0
    # i 番目の変数の条件（OK なら 0 を、NG なら 1 を返す）
    # i-1 番目までの変数に依存しても良い
    return 1 if (Z[i-1] + Z[i]) % 3 == 0 else 0

K = 10 # ループ変数の個数
Z = [0] * K
L = [0] * K
R = [0] * K
L[0], R[0] = 0, 1 # 半開区間 [L, R) をループする
i = 0
cnt = 0
while i >= 0:
    while i < K - 1:
        i += 1

        ##### 初期値・終了値の設定 #####
        L[i] = 0
        R[i] = sum(Z[:i]) + 2
        #####

        Z[i] = L[i]
        if chk(i):
            break
    else:
        ##### ここに処理を書く #####
        cnt += 1
        #####

    Z[i] += 1
    while Z[i] >= R[i] or chk(i):
        if Z[i] < R[i]: Z[i] += 1
        while Z[i] >= R[i]:
            i -= 1
            if i < 0: break
            Z[i] += 1
        if i < 0: break

print(cnt) # 5964602

ちゃんと同じ結果になりました。

問題（ネタバレ注意）

ABC 161 - D

問題
ループするだけですね。

ACコード
愚直ループ でもできるけど本番では書きたくないです。

ARC 095 - E

問題
想定解じゃないけど、自然に可変長ループすると枝刈りになって通ります 7 8。想定解の方法でも可変長ループが使えます。
ACコード

ARC 104 - E

問題
まあこれぐらいなら可変長ループいらないけど。
ACコード （古い書き方してたときのやつなのでこの記事の書き方とは違う）

Q&A

Q. これ需要あるの？
A. あんまりない気がしています。でも私は結構使ってます。直感的に書ける 9 ので気に入っています。

Q. 再帰じゃダメなの？
A. PyPy だと再帰遅いので 10 。もちろん計算時間に十分余裕があるなら大丈夫です。

Q. BFS じゃだめなの？
A. これは全然だめじゃないです。本記事の方法だとメモリが $O(K)$ で良いというメリットがありますが、メモリがネックになることはあまりない気がします。

Q. そもそも PyPy だと JIT コンパイルでめちゃくちゃメモリ食うから、メモリ $O(K)$ って意味なくない？
A. それを言われるととてもつらい気持ちになります。

Q. もっときれいに書けない？
A. 書けそう。教えてください。

Q. Step は設定できないの？（ループ変数を $2$ ずつ増やすとか）
A. ちょっといじればできます。

  1. 私が勝手に呼んでるだけですが ↩

  2. これができるものを「枝刈り可変長ループ」とか呼んだりしてます（私が勝手にry ↩

  3. 具体的にどう書くかはあんまり考えていません ↩

  4. この結果の 5964602 を覚えておいてね ↩

  5. 「さっきの例2」セクション参照 ↩

  6. ここでは一回ごとに $\Theta(k)$ をかけて設定していますが、ループしながら順に計算することで、それぞれ $O(1)$ で設定することもできます ↩

  7. AC コードでは可変長ループを二重に使っています。もちろん一重に全部入れることもできます ↩

  8. 嘘解法っぽいけど、嘘かどうかはこの記事の目的とは関係ないので気にしません ↩

  9. 個人の感想です ↩

  10. 再帰、ダメ、絶対 ↩

17
29
2

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
17
29