ABC129 D - Lamp

備忘録

問題

atcoder.jp

回答

import sys
import os
import math
import bisect
import itertools
import collections
import heapq
import queue
import array
import time
import datetime

# 時々使う
# import numpy as np
# from decimal import Decimal, ROUND_HALF_UP
# from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall

# 再帰の制限設定
sys.setrecursionlimit(10000000)


def ii(): return int(sys.stdin.buffer.readline().rstrip())
def il(): return list(map(int, sys.stdin.buffer.readline().split()))
def it(): return tuple(map(int, sys.stdin.buffer.readline().split()))
def fl(): return list(map(float, sys.stdin.buffer.readline().split()))
def iln(n): return [int(sys.stdin.buffer.readline().rstrip())
                    for _ in range(n)]


def iss(): return sys.stdin.buffer.readline().decode().rstrip()
def sl(): return list(map(str, sys.stdin.buffer.readline().decode().split()))
def isn(n): return [sys.stdin.buffer.readline().decode().rstrip()
                    for _ in range(n)]


def lcm(x, y): return (x * y) // math.gcd(x, y)


# MOD = 10 ** 9 + 7
MOD = 998244353
INF = float('inf')


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    H, W = il()
    grid = [iss() for _ in range(H)]

    # 予め前処理で、位置[i][j]のとき、
    # 左に何個照らすことが出来るか
    # 上に何個照らすことが出来るか求める
    L = [[0] * W for _ in range(H)]
    U = [[0] * W for _ in range(H)]
    for i in range(H):
        for j in range(W):
            if grid[i][j] == '#':
                continue
            L[i][j] = 1 if j == 0 else L[i][j-1] + 1
            U[i][j] = 1 if i == 0 else U[i-1][j] + 1

    # 予め前処理で、位置[i][j]のとき、
    # 右に何個照らすことが出来るか
    # 下に何個照らすことが出来るか求める
    R = [[0] * W for _ in range(H)]
    D = [[0] * W for _ in range(H)]
    for i in reversed(range(H)):
        for j in reversed(range(W)):
            if grid[i][j] == '#':
                continue
            R[i][j] = 1 if j == W - 1 else R[i][j+1] + 1
            D[i][j] = 1 if i == H - 1 else D[i+1][j] + 1

    # 前処理で得た値から位置[i][j]のとき
    # 上下左右に何個照らすことが出来るか求める
    ret = 0
    for i in range(H):
        for j in range(W):
            ret = max(ret, L[i][j] + R[i][j] + U[i][j] + D[i][j] - 3)
    print(ret)


if __name__ == '__main__':
    main()

考え方

計算量を減らす工夫と単純な実装力が問われる。
基本的には公式の解説の方針で、
ある位置Grid[i][j]を基準としたとき、
左右上下にいくつ照らすことができるかを累積和的に前処理で求める。

あとは全ての位置Grid[i][j]において、
あらかじめ求めた上下左右に照らすことが出来る数を合算すると回答となる。

処理で累積和を行う向きに注意が必要。
また、工夫を凝らさないとPython (3.8.2)では制限時間超過となる。
上記の回答例はPyPy3 (7.3.0)で提出。