AtCoder ABC 129 D - Lamp Pythonで解けなかったのでJuliaで(Python, Julia)

irisruneです。C問題ではテストケースでトラブルがありましたが、D問題ではPythonの実行時間で少し騒ぎになっていたようですね。

問題

atcoder.jp

公式解説と少しアルゴリズムが違いますが大体同じ解き方です。具体的には、連続する障害物のないマスの数を走査する際に左右の2方向、上下の2方向の結果を統合しているような感じです。

import copy
import sys
import numpy as np


def main():
    h, w = [int(x) for x in input().split()]
    board = [input() for _ in range(h)]
    boardint = [[1 if c == '.' else 0 for c in board[i]] for i in range(h)]
    ansrow = copy.deepcopy(boardint)
    for i in range(h):
        for j in range(1, w):
            ansrow[i][j] = ansrow[i][j - 1] + 1 if boardint[i][j] == 1 else 0
        for j in range(w - 2, -1, -1):
            ansrow[i][j] = ansrow[i][j + 1] if boardint[i][j:j+2] == [1, 1] else ansrow[i][j]
    anscol = copy.deepcopy(boardint)
    for j in range(w):
        for i in range(1, h):
            anscol[i][j] = anscol[i - 1][j] + 1 if boardint[i][j] == 1 else 0
        for i in range(h - 2, -1, -1):
            anscol[i][j] = anscol[i + 1][j] if boardint[i][j] == 1 and boardint[i + 1][j] == 1 else anscol[i][j]
    print(np.max(np.array(ansrow) + np.array(anscol)) - 1)


main()

（最後のnumpyを使う部分を修正した上で）PyPyで提出してもTLEでした。通ったコードと比較した感じ、入力したグリッドを1-0変換する部分やリストのコピーを行う部分が余計だったのかもしれません。

問題を解いている時には他人のコードを見ていなかったので、Pythonで通す方法がなかなか思いつきませんでした。そういうわけで文法の似ているJuliaで全く同じように書き直してみました。

parseInt(x)=parse(Int,x)
parseMap(x::Array{SubString{String},1})=map(parseInt,x)

function main()
    h, w = parseMap(split(readline()))
    board = zeros(Int, (h, w))
    for i in 1:h
            s = readline()
        for j in 1:w
            board[i, j] = s[j] == '.' ? 1 : 0
        end
    end
    ansrow = deepcopy(board)
    for i in 1:h
        for j in 2:w
            ansrow[i, j] = board[i, j] == 1 ? ansrow[i, j - 1] + 1 : 0
        end
        for j in w-1:-1:1
            ansrow[i, j] = board[i, j] == 1 && board[i, j + 1] == 1 ? ansrow[i, j + 1] : ansrow[i, j]
        end
    end
    anscol = deepcopy(board)
    for j in 1:w
        for i in 2:h
            anscol[i, j] = board[i, j] == 1 ? anscol[i - 1, j] + 1 : 0
        end
        for i in h-1:-1:1
            anscol[i, j] = board[i, j] == 1 && board[i + 1, j] == 1 ? anscol[i + 1, j] : anscol[i, j]
        end
    end
    print(maximum(ansrow + anscol) - 1)
end

main()

手元の環境（Julia1.1.0）とAtCoder上の環境（Julia0.5.0）との違いによりREを2回出していますが、とりあえず通りました。特に困ったのはrange()関数の仕様の違いでしたがコロンによる記法を使うことで解決しました。

PythonとJuliaのコードを見比べると一番目立つのはインデックスが0始まりか1始まりかという違いでしょうか。個人的には競プロの問題を解く上では1始まりの方がやりやすい気はしますね。Juliaでコードブロックを閉じるのに逐一endが必要なのは少々煩雑ですが仕方ないのかもしれません。

コードレビューとして、3項演算子を使うと（特に条件式が複合している場合に）1行が長くなりがちですね。JuliaはともかくPythonはどこに改行を挟んでよいかの判断が難しいです、リスト内法表記以外で3項演算子を使わない方がよいのかもしれませんが。

雑記

• 今PythonでAtCoderの問題を解いているのには一応理由があるのですが、今回のようなケースがあると色々考え直す必要があるかもしれないですね。
• Pythonで435msを叩き出している人がいるようで、わかる人が工夫すると割とどうにでもなるんだなと思いました。