はじめに
突然ですが、あなたは多重ループは書けますか？
？？？：「書けますー！」
cnt=0fori0inrange(3):fori1inrange(3):fori2inrange(3):fori3inrange(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$ を満たす整数を動くループを書いてください。
愚直に書くとこんな感じ。
cnt=0fori0inrange(1):fori1inrange(3):fori2inrange(5):fori3inrange(7):##### ここに処理を書く #####
cnt+=1#####
print(cnt)# 105
これぐらいなら Python の itertools を使うなどいろんな書き方があると思います3。
例2（愚直）
さっきのに加えて、「隣り合う変数の合計が $3$ の倍数になるものは除く」という条件を入れたらどうなるでしょうか。
cnt=0fori0inrange(1):fori1inrange(3):if(i0+i1)%3==0:continuefori2inrange(5):if(i1+i2)%3==0:continuefori3inrange(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)$ のように表される場合です。この例では右側だけ動かしていますが、左側が前の変数に依存することもありえます。
cnt=0fori0inrange(1):fori1inrange(i0+2):if(i0+i1)%3==0:continuefori2inrange(i0+i1+2):if(i1+i2)%3==0:continuefori3inrange(i0+i1+i2+2):if(i2+i3)%3==0:continue##### ここに処理を書く #####
cnt+=1#####
print(cnt)# 5
これのループ変数の個数を増やすと（例えば $K=10$ など）結構大変そうですね。
まあ書けなくはないんですけど。
cnt=0fori0inrange(1):fori1inrange(i0+2):if(i0+i1)%3==0:continuefori2inrange(i0+i1+2):if(i1+i2)%3==0:continuefori3inrange(i0+i1+i2+2):if(i2+i3)%3==0:continuefori4inrange(i0+i1+i2+i3+2):if(i3+i4)%3==0:continuefori5inrange(i0+i1+i2+i3+i4+2):if(i4+i5)%3==0:continuefori6inrange(i0+i1+i2+i3+i4+i5+2):if(i5+i6)%3==0:continuefori7inrange(i0+i1+i2+i3+i4+i5+i6+2):if(i6+i7)%3==0:continuefori8inrange(i0+i1+i2+i3+i4+i5+i6+i7+2):if(i7+i8)%3==0:continuefori9inrange(i0+i1+i2+i3+i4+i5+i6+i7+i8+2):if(i8+i9)%3==0:continue##### ここに処理を書く #####
cnt+=1#####
print(cnt)# 5964602
これをきれいに書くのを目標にしましょう 4。
他の方法
本記事で示す方法以外にもいくつか方法があります。これらが使えることもあるでしょう。
愚直にループを書く
上の例で書いたようなやつです。ただこれだとループの深さなどが可変の問題には対応しづらいですね。itertools
不要なところもループしてしまうので場合によっては計算量が増えてしまう可能性があります。再帰 DFS
再帰は言語によっては遅くなるので使いたくなかったりします。BFS
メモリが $O(K)$ では収まらないです。本記事の方法
可変長対応、不要なループをしない 5、非再帰、メモリは $O(K)$ 、のすべてを満足する方法です。アイデア
アイデアというほどでもないですが、本記事の実装方法です。
- 各ループ変数に対して、「初期値」「終了値」を持つ。具体的には $i$ 番目のループ変数の初期値・終了値をそれぞれ
L[i]およびR[i]で表す - while 文をでループする
これだけです。
実装
さっきの例1
さっきの例で実装してみます。chk関数で満たすべき条件を設定しています。この例では特に NG 条件がないので必ず $0$ を返すようにしています。
defchk(i):ifnoti:return0# i 番目の変数の条件（OK なら 0 を、NG なら 1 を返す）
# i-1 番目までの変数に依存しても良い
return0K=4# ループ変数の個数
Z=[0]*KL=[0]*KR=[0]*KL[0],R[0]=0,1# 半開区間 [L, R) をループする
i=0cnt=0whilei>=0:whilei<K-1:i+=1##### 初期値・終了値の設定 #####
L[i]=0R[i]=2*i+1#####
Z[i]=L[i]ifchk(i):breakelse:##### ここに処理を書く #####
cnt+=1#####
Z[i]+=1whileZ[i]>=R[i]orchk(i):ifZ[i]<R[i]:Z[i]+=1whileZ[i]>=R[i]:i-=1ifi<0:breakZ[i]+=1ifi<0:breakprint(cnt)# 105
さっきの例2
ループ変数間に条件がある場合です。chk関数に条件を入れています。 $i$ 番目のループ変数の条件のところでは $i$ 番目までの変数に依存して決めて良いです。言い換えると、いくつかのループ変数に関係する条件については、その中で最も大きな番号のループ変数のところで判定をして、 NG ならすぐ次に移るようにしています。
defchk(i):ifnoti:return0# i 番目の変数の条件（OK なら 0 を、NG なら 1 を返す）
# i-1 番目までの変数に依存しても良い
return1if(Z[i-1]+Z[i])%3==0else0K=4# ループ変数の個数
Z=[0]*KL=[0]*KR=[0]*KL[0],R[0]=0,1# 半開区間 [L, R) をループする
i=0cnt=0whilei>=0:whilei<K-1:i+=1##### 初期値・終了値の設定 #####
L[i]=0R[i]=2*i+1#####
Z[i]=L[i]ifchk(i):breakelse:##### ここに処理を書く #####
cnt+=1#####
Z[i]+=1whileZ[i]>=R[i]orchk(i):ifZ[i]<R[i]:Z[i]+=1whileZ[i]>=R[i]:i-=1ifi<0:breakZ[i]+=1ifi<0:breakprint(cnt)# 31
さっきの例3
$i$ 番目のループ変数の範囲が $i-1$ 番目までのループ変数に依存する場合です。R[i] = sum(Z[:i]) + 2のところで設定しています 6。
defchk(i):ifnoti:return0# i 番目の変数の条件（OK なら 0 を、NG なら 1 を返す）
# i-1 番目までの変数に依存しても良い
return1if(Z[i-1]+Z[i])%3==0else0K=4# ループ変数の個数
Z=[0]*KL=[0]*KR=[0]*KL[0],R[0]=0,1# 半開区間 [L, R) をループする
i=0cnt=0whilei>=0:whilei<K-1:i+=1##### 初期値・終了値の設定 #####
L[i]=0R[i]=sum(Z[:i])+2#####
Z[i]=L[i]ifchk(i):breakelse:##### ここに処理を書く #####
cnt+=1#####
Z[i]+=1whileZ[i]>=R[i]orchk(i):ifZ[i]<R[i]:Z[i]+=1whileZ[i]>=R[i]:i-=1ifi<0:breakZ[i]+=1ifi<0:breakprint(cnt)# 5
さっきみたいに $K=10$ にしてみましょう。
defchk(i):ifnoti:return0# i 番目の変数の条件（OK なら 0 を、NG なら 1 を返す）
# i-1 番目までの変数に依存しても良い
return1if(Z[i-1]+Z[i])%3==0else0K=10# ループ変数の個数
Z=[0]*KL=[0]*KR=[0]*KL[0],R[0]=0,1# 半開区間 [L, R) をループする
i=0cnt=0whilei>=0:whilei<K-1:i+=1##### 初期値・終了値の設定 #####
L[i]=0R[i]=sum(Z[:i])+2#####
Z[i]=L[i]ifchk(i):breakelse:##### ここに処理を書く #####
cnt+=1#####
Z[i]+=1whileZ[i]>=R[i]orchk(i):ifZ[i]<R[i]:Z[i]+=1whileZ[i]>=R[i]:i-=1ifi<0:breakZ[i]+=1ifi<0:breakprint(cnt)# 5964602
ちゃんと同じ結果になりました。
問題（ネタバレ注意）
ABC 161 - D
問題
ループするだけですね。
ACコード
愚直ループでもできるけど本番では書きたくないです。
ARC 095 - E
問題
想定解じゃないけど、自然に可変長ループすると枝刈りになって通ります 78。想定解の方法でも可変長ループが使えます。
ACコード
ARC 104 - E
問題
まあこれぐらいなら可変長ループいらないけど。
ACコード（古い書き方してたときのやつなのでこの記事の書き方とは違う）
Q&A
Q. これ需要あるの？
A. あんまりない気がしています。でも私は結構使ってます。直感的に書ける 9ので気に入っています。
書いても需要ない記事ばかり書く人というのが存在していてほしい
— えびちゃん (@rsk0315_h4x) March 3, 2021
Q. 再帰じゃダメなの？
A. PyPy だと再帰遅いので 10。もちろん計算時間に十分余裕があるなら大丈夫です。
Q. BFS じゃだめなの？
A. これは全然だめじゃないです。本記事の方法だとメモリが $O(K)$ で良いというメリットがありますが、メモリがネックになることはあまりない気がします。
Q. そもそも PyPy だと JIT コンパイルでめちゃくちゃメモリ食うから、メモリ $O(K)$ って意味なくない？
A. それを言われるととてもつらい気持ちになります。
Q. もっときれいに書けない？
A. 書けそう。教えてください。
Q. Step は設定できないの？（ループ変数を $2$ ずつ増やすとか）
A. ちょっといじればできます。