Polaris.AI プログラミングコンテスト 2025（AtCoder Beginner Contest 429）

次の言語アップデートで導入されるCodonについて調べていた矢先、早速、打って付けの問題が来た。
アップデート後に来てくれていれば。。。

F - Shortest Path Query

問題文

$3$ 行 $N$ 列のグリッドが与えられます。# ならば壁マスで、 . ならば空きマスであり通行可能です。
$Q$ 個のクエリが与えられるので、順に処理してください。
- 各クエリでは整数 $r,c$ が与えられるので、マス $(r,c)$ の状態を反転させてください。
  - つまり、マス $(r,c)$ が壁マスならば空きマスにし、空きマスならば壁マスにしてください。
- その後、以下の問題の答えを出力してください。
  - マス $(1,1)$ から上下左右に隣接する空きマスに移動する操作を繰り返してマス $(3,N)$ に移動することを考えます。
  - このとき、マス $(3,N)$ に到達できるか判定し、到達できる場合は操作回数の最小値を求めてください。

制約

$2\le N\le 2\times 10^5$
$S_{i,j}$ は # または .
$S_{1,1}=S_{3,N}=$ .
$1\le Q\le 2\times 10^5$
$1\le r\le 3$
$1\le c\le N$
$(r,c) \neq (1,1),(3,N)$
$N,Q,r,c$ は整数

解法

グリッド上の経路で「左に戻る」のが最短になるには、下のように最低5行必要になる。
3行しかないこの問題では、最短経路において左移動は発生しない。

......
#####.
......
.#####
......

つまり、セグメント木に「ある区間の最左○行目から最右○行目までの最短距離」を載せると、マージの際は「左区間の右端→右区間の左端間の移動を、何行目でおこなうか」だけ気にすればよくなる。

■セグメント木に載せる情報
  0: 単位元かどうか?
  1: 最左1行目→最右1行目  2: 最左1行目→最右2行目  3: 最左1行目→最右3行目
  4: 最左2行目→最右1行目  5: 最左2行目→最右2行目  6: 最左1行目→最右3行目
  7: 最左3行目→最右1行目  8: 最左3行目→最右2行目  9: 最左3行目→最右3行目
     ...のそれぞれの最短距離（到達不能や、左右いずれかの該当マスが壁ならINF）

マージは、いずれかが単位元ならもう一方がそのまま採用される。

どちらも意味のある区間なら、左右の間で3行のどこを通過したかを全て試す。
通過する際に $+1$ が発生する点に注意。

セグメント木に載せる情報の idx:1～9 の $9$ 個を埋めるのに $27$ 回の計算で可能となる。定数倍が重め。

         +1           たとえばマージ後の idx:1（1行目→1行目）を求めたければ、
....##.. → .##...##    左1→1 → 1→1右
###..#.# → ..#.#.#.    左1→2 → 2→1右  の3通りを試し、最小値をとればよい。
...#.... → ....#...    左1→3 → 3→1右

高速化について

こういう、持たせる情報の数が多くなるセグメント木は、特にPython,PyPyは遅くなりがち。
なかなかTLEが取れず、Numbaで実装しなおしてようやく通った。

「Python で class を使ってライブラリ化しているデータ構造」は、 Numbaはいまいち使いづらい（jitclass はあるが、結局書き換えが必要だし、まだexperimental）ので、書き換えになかなか踏ん切りがつきづらい。

なお Codon だと、細かな点でエラーが出たが（@classmethod や int.bit_length() が使えない）、これらは少ない書き換え量で対応でき、コンパイルを通すことができた。class もそのまま使えるのは大きい。

それでも Python で通している人はいるもので、テスターや、コンテスト中にACしている他の方の提出を見ると、だいたいTLEした自分の方針と大きくは違わないようには見えるのだが、以下の傾向がある？

タプルではなくリストで持たせる（参照渡しにならないよう注意）
マージ時の各最小値の計算は、展開せず for 文で回す

タプルよりリストの方が、わずかながら生成が高速っぽい。

また、意外なのが、マージ時に各最短距離を求めるのに

return [
    min(a[1]+b[1], a[2]+b[4], a[3]+b[7]) + 1,
    min(a[1]+b[2], a[2]+b[5], a[3]+b[8]) + 1,
    ...
]

より、

c = [INF]*9
for i in range(3):
    for j in range(3):
        for k in range(3):
            c[i*3+j] = min(c[i*3+j], a[i*3+k]+b[k*3+j]+1)

のような実装の方が、明確にPyPyで速くなった。

リストアクセスや、+1 の演算回数など、前者の方が少ないはずなんだけど。
やっぱりPyPyの挙動、非自明すぎる。。。

Python3

import sys

import numpy as np


def solve_numba(n, states, queries):
    INF = 1 << 60

    # == SegmentTree ==
    # Numbaで使うにあたり、いくつかのアドホックな使用方法に注意する。
    # ・data は2次元配列 (index, 1要素当たりのデータ数) で、Numbaで扱える型。型は配列内で全て揃える必要がある。
    #     末尾2つをクエリ計算中の値の保持に使うため、2だけ長く作る。
    # ・使用前に以下の関数を定義する。(prefixの"sgt_"は省略)
    #     operation(i,j,k):  data[k] に data[i]+data[j] を上書き更新（k=i や k=j であることもあるため注意）
    #     initialize_data(i): data[i] を単位元に戻す
    # ・初期配列があるときでも、自動的に構築はされない。
    #   data[sgt_offset]～ 以降に値をセットした後、葉に近い方から sgt_update(i) を呼んでいって自前で前計算しておく

    sgt_n = n  # CUSTOMIZE
    sgt_depth = 0
    _n = sgt_n
    while _n:
        _n >>= 1
        sgt_depth += 1
    sgt_offset = 1 << sgt_depth

    # ==== CUSTOMIZE ====
    sgt_data = np.zeros(((sgt_offset << 1) + 2, 10), np.int64)

    def sgt_operation(i, j, k):
        if sgt_data[i, 0] == 1:
            if j != k:
                sgt_data[k] = sgt_data[j, :]
            return
        if sgt_data[j, 0] == 1:
            if i != k:
                sgt_data[k] = sgt_data[i, :]
            return
        a0, a1, a2, a3, a4, a5, a6, a7, a8, a9 = sgt_data[i]
        b0, b1, b2, b3, b4, b5, b6, b7, b8, b9 = sgt_data[j]
        sgt_data[k, 0] = 0
        sgt_data[k, 1] = min(INF, a1 + b1, a2 + b4, a3 + b7) + 1
        sgt_data[k, 2] = min(INF, a1 + b2, a2 + b5, a3 + b8) + 1
        sgt_data[k, 3] = min(INF, a1 + b3, a2 + b6, a3 + b9) + 1
        sgt_data[k, 4] = min(INF, a4 + b1, a5 + b4, a6 + b7) + 1
        sgt_data[k, 5] = min(INF, a4 + b2, a5 + b5, a6 + b8) + 1
        sgt_data[k, 6] = min(INF, a4 + b3, a5 + b6, a6 + b9) + 1
        sgt_data[k, 7] = min(INF, a7 + b1, a8 + b4, a9 + b7) + 1
        sgt_data[k, 8] = min(INF, a7 + b2, a8 + b5, a9 + b8) + 1
        sgt_data[k, 9] = min(INF, a7 + b3, a8 + b6, a9 + b9) + 1

    def sgt_initialize_data(i):
        sgt_data[i, 0] = 1

    # ==== CUSTOMIZE ここまで。以降の sgt_ から始まる関数は基本的には触らなくていい ====

    def sgt_update(i):
        lch = i << 1
        rch = lch + 1
        sgt_operation(lch, rch, i)

    def sgt_point_set(index, value):
        p = index + sgt_offset
        sgt_data[p] = value
        for i in range(1, sgt_depth + 1):
            sgt_update(p >> i)

    def sgt_point_get(index):
        p = index + sgt_offset
        return sgt_data[p]

    def sgt_range_query(l, r):
        l += sgt_offset
        r += sgt_offset
        sml = sgt_offset << 1
        smr = sml + 1
        sgt_initialize_data(sml)
        sgt_initialize_data(smr)

        while l < r:
            if (l & 1) == 1:
                sgt_operation(sml, l, sml)
                l += 1
            if (r & 1) == 1:
                r -= 1
                sgt_operation(r, smr, smr)
            l >>= 1
            r >>= 1
        sgt_operation(sml, smr, sml)
        return sgt_data[sml]

    def sgt_debug_print():
        i = 1
        while i <= sgt_offset:
            print(sgt_data[i:i * 2])
            i <<= 1

    INF = 1 << 60
    state_dic = [
        (0, 0, 1, 2, 1, 0, 1, 2, 1, 0),
        (0, INF, INF, INF, INF, 0, 1, INF, 1, 0),
        (0, 0, INF, INF, INF, INF, INF, INF, INF, 0),
        (0, INF, INF, INF, INF, INF, INF, INF, INF, 0),
        (0, 0, 1, INF, 1, 0, INF, INF, INF, INF),
        (0, INF, INF, INF, INF, 0, INF, INF, INF, INF),
        (0, 0, INF, INF, INF, INF, INF, INF, INF, INF),
        (0, INF, INF, INF, INF, INF, INF, INF, INF, INF),
    ]

    # (0→0):1 (0→1):2 (0→2):3
    # (1→0):4 (1→1):5 (1→2):6
    # (2→0):7 (2→1):8 (2→2):9

    for i in range(n):
        sgt_data[i + sgt_offset, :] = state_dic[states[i]]
    for i in range(n + sgt_offset, sgt_offset * 2):
        sgt_data[i, 0] = 1
    for i in range(sgt_offset - 1, 0, -1):
        sgt_update(i)

    buf = []
    for r, c in queries:
        r -= 1
        c -= 1
        states[c] ^= 1 << r
        sgt_point_set(c, state_dic[states[c]])
        ans = sgt_data[1, 3]
        if ans >= INF:
            buf.append(-1)
        else:
            buf.append(ans)

    return buf


def solve(n, grid, q, queries):
    states = np.zeros(n, np.int64)
    for i in range(n):
        states[i] = (grid[0][i] == '#') | ((grid[1][i] == '#') << 1) | ((grid[2][i] == '#') << 2)
    queries = np.array(queries, np.int64)
    ans = solve_numba(n, states, queries)
    return ans


SIGNATURE = '(i8,i8[:],i8[:,:])'
# Compile Phase
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve_numba', SIGNATURE)(solve_numba)
    cc.compile()
    exit()

# Execute Phase
if len(sys.argv) >= 3:
    # noinspection PyUnresolvedReferences
    from my_module import solve_numba
else:
    from numba import njit

    solve_numba = njit(SIGNATURE, cache=True)(solve_numba)
    print('compiled', file=sys.stderr)

n = int(input())
grid = [input() for _ in range(3)]
q = int(input())
queries = [list(map(int, input().split())) for _ in range(q)]
ans = solve(n, grid, q, queries)
print('\n'.join(map(str, ans)))

目次

Polaris.AI プログラミングコンテスト 2025（AtCoder Beginner Contest 429）

F - Shortest Path Query

問題文

制約

解法

高速化について