ベスパリブ

プログラミングを主とした日記・備忘録です。ベスパ持ってないです。

ABC 129: D問題の解説 (Python3)

問題: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()

計算量は、 O(HW)

単純な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]