グリッドを渦巻き状に探索したい

きっかけ

ABC335: D - Loong and Takahashiでグリッドを渦巻き状に探索する必要があった.

やりたいこと

以下の図のように, グリッドの左上から初めて渦巻き状にグリッドの中心まで探索したい.

虚数の回転を使った解法

基本的に探索する方向は,

1. →
2. ↓
3. ←
4. ↑

の順番を繰り返すので, 以下のように現在の探索している方向directionを虚数として管理する (座標軸のy方向とgrid上のy方向が逆転していることに注意).

1. → = 0 + 1j
2. ↓ = 1 + 0j
3. ← = 0 - 1j
4. ↑ = -1 + 0j

1 -> 2, 2 -> 3, 3 -> 4, 4 -> 1 と方向を変えるには, -1j を かければ良い.

同様に現在のマス目も虚数nowで管理することで, 現在の地点から次の地点への移動は, now + direction で求めることができる.

コード

最終的に以下のような感じになる.

N = int(input())

grid = [[""] * N for _ in range(N)]
grid[N // 2][N // 2] = "T"

# h + wj
now = 0 + 0j
cnt = 1
direction = 0 + 1j

while grid[int(now.real)][int(now.imag)] != "T":
    grid[int(now.real)][int(now.imag)] = str(cnt)

    # 方向を変える必要があるのは, 既に探索済 or gridから飛び出るとき
    nh, nw = int(now.real + direction.real), int(now.imag + direction.imag)
    if (nh < 0) or (N <= nh) or (nw < 0) or (N <= nw) or (grid[nh][nw] != ""):
        direction *= -1j

    now += direction
    cnt += 1

for g in grid:
    print(*g)

Verify