ABC335・D - Loong and Takahashi

問題はこちら

AtCoder Beginner Contest 335（Sponsored by Mynavi）
D問題

提出した回答

n = int(input())

# 2次元配列を生成
grid = [[0 for i in range(n)] for j in range(n)]
grid[int((n+1)/2-1)][int((n+1)/2-1)] = "T"

x, y = 0, 0
cnt = 0
    
direction = ["right", "down", "left", "up"]

def go(grid, direction, cnt, x, y):
    if direction[cnt] == "right" and x >= 0 and y+1 >= 0 and grid[x][y+1] == 0:
        y += 1
        return x, y
    elif direction[cnt] == "left" and x >= 0 and y-1 >= 0 and grid[x][y-1] == 0:
        y -= 1
        return x, y
    elif direction[cnt] == "up" and x-1 >= 0 and y >= 0 and grid[x-1][y] == 0:
        x -= 1
        return x, y
    elif direction[cnt] == "down" and x+1 >= 0 and y >= 0 and grid[x+1][y] == 0:
        x += 1
        return x, y
    else:
        return False

for i in range(n**2 - 1):
    if i == 0:
        grid[x][y] = i + 1
    else:
        try:
          x, y = go(grid, direction, cnt, x, y)
          grid[x][y] = i + 1
    except:
        if cnt != 3:
            cnt += 1
        else:
            cnt = 0
          x, y = go(grid, direction, cnt, x, y)
          grid[x][y] = i + 1

for i in range(len(grid)):
    print(*grid[i])

苦戦したところ

処理としては、初めに一番左上に1を入力し、ひたすら右へ。
埋められなくなったら下へ、また埋められなくなったら今度は左へ……と、「今の処理ができなくなるまで続け、不可能になったら処理を切り替える」というプログラムを作成しました。
この、「今の処理ができなくなるまで続け、不可能になったら処理を切り替える」をどのように実装するかがすぐに分からず、苦戦しました。

実装方法

今回は、
①右へ
②下へ
③左へ
④上へ
を順に処理していく、というプログラムです。
なので、

direction = ["right", "down", "left", "up"]

というリストを作成し、

def go(grid, direction, cnt, x, y):
    if direction[cnt] == "right" and x >= 0 and y+1 >= 0 and grid[x][y+1] == 0:
        y += 1
        return x, y
    elif direction[cnt] == "left" and x >= 0 and y-1 >= 0 and grid[x][y-1] == 0:
        y -= 1
        return x, y
    elif direction[cnt] == "up" and x-1 >= 0 and y >= 0 and grid[x-1][y] == 0:
        x -= 1
        return x, y
    elif direction[cnt] == "down" and x+1 >= 0 and y >= 0 and grid[x+1][y] == 0:
        x += 1
        return x, y
    else:
        return False

のように、リストのインデックスを指定することで処理を切り分ける関数を作成しました。
そして、for文の中で、

try:
    x, y = go(grid, direction, cnt, x, y)
    grid[x][y] = i + 1
except:
    if cnt != 3:
        cnt += 1
    else:
        cnt = 0
    x, y = go(grid, direction, cnt, x, y)
    grid[x][y] = i + 1

と、関数がエラーを返したら処理内容リストのインデックスであるcntをインクリメントして再実行する、としています。


とりあえずこれで正解はもらえましたが、もっと良い方法があると思うので調べてみます。