ABC311、F問題の初心者的解説

0. はじめに

AtCoder Beginner Contest 311（ABC311、トヨタ自動車プログラミングコンテスト2023#4）のF問題について、解説を読んでもよく分かりませんでした。初心者的に樹形図的に考えてなんとかACしたので、自分で解説を書いてみました。

本記事の内容は、つぎのとおりです。

1. 問題設定、「美しいグリッド」について
2. 樹形図による数え上げ（辞書型で実装）
3. テーブル中の累積和による高速化
4. 終わりに

1. 問題設定、「美しいグリッド」について

美しいグリッドの定義からは、解法に必要な以下の二点が導かれます。

• 各列の最も上にある「#」の位置が必要。
• 各列を左から右（$j$の増加方向）に調べる。

問題文の定義から、美しいグリッドは直角二等辺三角形の近似表現である事が分かります。ある一点 $(i, j)$ が「#」ならば、その下 $(i+1, j)$ と右下 $(i+1, j+1)$ が「#」となり、これが下方（$i$増加方向）へ延々と続きます。この直角二等辺三角形の形状、加えてそれが複数重なってできる鋸歯状の形状が数え上げの制限条件となります。

例えば、初期入力として5x5のグリッドの位置 $(1, 1)$ のみに「#」が与えられた場合、初期状態は以下のようになります。

. . . . .
. # . . .
. ## . .
. ### .
. #####

求められているのは「美しいグリッドの種類数」なので、入力値をまずは美しいグリッドに直す必要があります。これが初期状態です。二重のfor文によってこの初期状態を得ることは可能ですが、実は問題を解く上ではこの初期状態のグリッド全体の情報は必要ありません。必要なのは、各列の「#」の最も上の位置の情報です。

なぜなら、「美しいグリッドの種類数」のために必要な情報は、初期状態から新たに塗りつぶし可能な空白マスの個数（位置）情報だからです。美しいグリッドの定義より各列の「#」の下は全て「#」なので、空白マスの個数の情報は、各列の「#」の最も上の位置の情報と同等となります。

従って解法に必要な初期状態は、入力値と美しいグリッドの定義に従った、各列の「#」の最上位の位置情報です。

もう一つ、美しいグリッドの定義から導かれることは、列を左から右に調べることです。ある列の位置 $(i, j)$ に「#」が存在すれば、右隣の列の位置 $(i+1, j+1)$ に「#」が存在するからです。従って、左の列 $(j=0)$ から順に状態の数を数え上げていくことで、「美しいグリッドの種類数」を数え上げることができます。

2. 樹形図による数え上げ（辞書型で実装）

まず辞書型で実装したコードを紹介します。このコードでは、いくつかのケースがTLEになります。

abc311_f_dict.py

from collections import defaultdict

mod = 998244353
n, m = map(int, input().split())
p = [list(input()) for _ in range(n)]

# 初期状態p1の導出
# 各列の0<=i<=n-1に「#」が存在しない場合、「#」がi=nにあると設定、
p1=[n]*m
for j in range(m):
  for i in range(n):
    if p[i][j]=="#":
      p1[j]=i
      break
  if j!=0:
    p1[j] = min(p1[j-1]+1, p1[j])

# 辞書による樹形図の実装
d = [defaultdict(int)]

# 最初に、j=0の各空白マスを塗りつぶした新たな美しいグリッド状態として１を与える
# key = マスの位置
# val = 状態数
# 「＃」の位置の状態数１は、「新たな塗りつぶしをしない」選択を意味する
# すでに塗りつぶされている位置の選択は、「新たな塗りつぶしをしない」選択と同義
for i in range(p1[0]+1):
  d[-1][i]=1
  
for j in range(1,m):
  d.append(defaultdict(int))
# 左隣の列の各マスの状態を呼び出し。
  for key, val in d[-2].items():
# 塗りつぶしによって得られた新たな美しいグリッドと初期条件を比較し、
# 選択可能なマスに状態数を加える
# 先程同様、「新たな塗りつぶしをしない」選択にも状態数を加える
    key2 = min(key+1, p1[j])
    for k in range(key2+1):
      d[-1][k] += val
  for key, val in d[-1].items():
    d[-1][key] = val%mod
    
ans = 0
for key, val in d[-1].items():
  ans = (ans+val)%mod
print(ans)

先に書いたように、今回の問題では左の列から調べれば良いことが分かっています。そのために樹形図的な数え上げの方法として、最も左端の$j=0$の列の各マスに初期の状態数１を与えて、右隣の列への配分を計算します。

この状態の配分について考えます。例えば$(i,0)$を塗りつぶして美しいグリッドを作った後で、初期状態の邪魔を無視すると、右隣の$(0, 1)$から$(i, 1)$の一マスを次に塗りつぶして種類の異なる美しいグリッドを作ることができます。これを実装しようとすると、$(i,0)$の状態数1を、右隣の$(0, 1)$から$(i+1, 1)$のマスに加算することになります。範囲の一つ広がった$(i+1, 1)$は、「その列を新たに塗りつぶさない」選択です。すでに美しいグリッドで塗りつぶされているためです。

このようにこの右隣の列への配分の際、「美しいグリッド」という制限があるため配分可能な右隣のマスに制限ができます。これがコード中の「key2 = min(key+1, p1[j])」です。

ただ、このコードはいくつかのケースでTLEになります。原因は、数え上げの際の三重ループです。

for j in range(1,m):
  for key, val in d[-2].items():
    for k in range(key2+1):

ある列の選択可能なマスの数は最大で$N$なので、コード中のkeyとkey2の最大選択範囲も$N$であり、この三重ループの計算量は$MN^2$となります。この計算量を回避する方法が、累積和の使用です。三重目のループの中身が単純な加算（d[-1][k] += val）なので累積和が可能です。

3. テーブル中の累積和による高速化

このコードの変更点は、次のとおりです。

• 辞書によるメモ化の代わりにテーブルを使用。
• 累積和の使用。

abc311_f_table.py

mod = 998244353
n, m = map(int, input().split())
p = [list(input()) for _ in range(n)]

# 初期状態
p1=[n]*m
for j in range(m):
  for i in range(n):
    if p[i][j]=="#":
      p1[j]=i
      break
  if j!=0:
    p1[j] = min(p1[j-1]+1, p1[j])

# テーブルの使用
dp = [[0]*m for _ in range(n+3)]
for i in range(n+1):
  dp[i][0]=1

for j in range(1,m):
# 累積和の計算
  for i in range(p1[j-1],-1,-1):
    dp[i+1][j] = dp[i+2][j] + dp[i][j-1]
  dp[0][j] = dp[1][j]
  for i in range(p1[j-1]+3):
    dp[i][j] %= mod
    
ans = 0
j=m-1
for i in range(p1[j]+1):
  ans = (ans+dp[i][j])%mod
print(ans)

累積和の使用によって三重ループを二重ループにすることができました。これによって計算量は$MN$となり、ACしました。

テーブルの変数名を「dp」としましたが、これをDPとみなして良いのでしょうか？みなすなら「dp[i+1][j] = dp[i+2][j] + dp[i][j-1]」が累積和を使用した漸化式となります。

4. 終わりに

公式の解説を読んでも「ん～？」と唸るしかなかったので、自分で初心者？用の解説を書いてみました。今回紹介したコードが、よく練られたDPの解法とは異なることをご承知おき下さい。

樹形図と辞書を使用した動的計画法の解法については、こちらの記事を御覧ください。

• 動的計画法の辞書型データ構造による実装、