問題:AtCoder Beginner Contest 129: D - Lamp
公式解説PDF:https://img.atcoder.jp/abc129/editorial.pdf
公式解説動画:https://youtu.be/L8grWxBlIZ4?t=4363
ABC129のD問題のPython3による解説です。
解説
DP問題です。
これは公式の解説動画がとてもわかりやすかったので、そちらを見るのをおすすめします。
問題を解くアイデアとしては、i行j列目に明かりを設置したとき、
- 左方向に照らせるマスの数を格納した L
- 右方向に照らせるマスの数を格納した R
- 上方向に照らせるマスの数を格納した U
- 下方向に照らせるマスの数を格納した D
のそれぞれを作成します。L, R, U, DはシンプルなDPによって計算することができます。
すると、i行j列目に明かりを設置したとき、上下左右に明かりを照らせるマスの数を格納したものを ans_grid とすると、
ans_grid[i][j] = L[i][j] + R[i][j] + U[i][j] + D[i][j] - 3
で導出することができるので、あとは ans_grid の要素の最大値を出力すればOKです。
コーナーケース
- 特になし
実装
# -*- coding:utf-8 -*- import sys import numpy as np def solve(): H, W = list(map(int, sys.stdin.readline().split())) grid = [] for _ in range(H): grid.append(list(input())) grid = (np.array(grid) == ".")*1 # ランプが置ける:1, 障害物:0 #print(grid) L = np.zeros((H,W), dtype=int) # 左方向に照らせるマスの数 R = np.zeros((H,W), dtype=int) # 右方向に照らせるマスの数 U = np.zeros((H,W), dtype=int) # 上方向に照らせるマスの数 D = np.zeros((H,W), dtype=int) # 下方向に照らせるマスの数 # 左方向に照らせるマスの数のDPを埋める for i in range(W): if i == 0: L[:, i] = grid[:, i] else: L[:, i] = (L[:, i-1] + 1) * grid[:, i] # 右方向に照らせるマスの数のDPを埋める for i in range(W-1, -1, -1): if i == W-1: R[:, i] = grid[:, i] else: R[:, i] = (R[:, i+1] + 1) * grid[:, i] # 上方向に照らせるマスの数のDPを埋める for i in range(H): if i == 0: U[i] = grid[i] else: U[i] = (U[i-1] + 1) * grid[i] # 下の方向に照らせるマスの数のDPを埋める for i in range(H-1, -1, -1): if i == H-1: D[i] = grid[i] else: D[i] = (D[i+1] + 1) * grid[i] # 上下左右すべてのDPを埋めた。 # よって、i行j列目に明かりを設置したとき、上下左右に照らせるマスの数を格納した ans_grid は、 ans_grid = L + R + U + D - 3 ans = np.max(ans_grid) print(ans) if __name__ == "__main__": solve()
計算量は、
単純なforループでやろうとするとTLEしたので、行列計算を高速でしてくれるnumpyを使っています。
numpyの重要な特徴
narrayに対する演算
L = np.zeros((4,3), dtype=int) # 全ての要素が0の、4行3列の行列を作成 print(L+1) # [[1 1 1] # [1 1 1] # [1 1 1] # [1 1 1]]
行を順番に取り出す
R = np.arange(4*3).reshape(4, 3) # 4行3列の行列を作成 print(R) # [[ 0 1 2] # [ 3 4 5] # [ 6 7 8] # [ 9 10 11]] # 行を順番に取り出す for i in range(4): print(R[i]) # [0 1 2] # [3 4 5] # [6 7 8] # [ 9 10 11]
列を順番に取り出す
for i in range(3): print(R[:, i]) # [0 3 6 9] # [ 1 4 7 10] # [ 2 5 8 11]