[ABC258F] Main Street を Pythonで解く

https://atcoder.jp/contests/abc258/tasks/abc258_f

考え方

$(S_x, S_y)$ から 高速道路・普通道を使いながら $(G_x, G_y)$ へ移動する問題です。高速道路の方が所要時間が短いので、できるだけ高速道路を使う方が良いです。$(S_x, S_y)$ と $(G_x, G_y)$ から4方向へ移動して高速道路と交わる地点が高速道路の乗り口・降り口の候補となります。$(S_x, S_y)$ と $(G_x, G_y)$ がそれぞれ4点ずつ候補を持つため、4×4=16通りを探索する必要があります。

実装メモ

$(x,y)$ を4方向へ移動したときに、高速道路と交わる地点の座標を求める関数を準備します。同時に高速道路に到達する前に（降りた後に）普通道を通る距離も計算しておきます。

def highway_pos(B,x,y):
    x1,y1,x2,y2 = x//B*B, y//B*B, (x//B+1)*B, (y//B+1)*B
    dx,dy = x%B, y%B
    pos = [[x1,y,dx], [x2,y,B-dx], [x,y1,dy], [x,y2,B-dy]]
    return pos

高速道路上の地点 $(x_1, y_1)$ と $(x_2, y_2)$ の距離は、ほとんどの場合は、

$$abs(x_1-x_2) + abs(y_1-y_2)$$

で求めることができますが、 $(x_1, y_1)$ と $(x_2, y_2)$ が上記の式とはならない場合があります。

def highway_dist(B,x1,y1,x2,y2):
    if x1%B==0 and x2%B==0 and y1//B==y2//B and x1//B!=x2//B:
        return min(y1%B+y2%B, B*2-y1%B-y2%B) + abs(x1-x2)

    if y1%B==0 and y2%B==0 and x1//B==x2//B and y1//B!=y2//B:
        return min(x1%B+x2%B, B*2-x1%B-x2%B) + abs(y1-y2)
    
    return abs(x1-x2) + abs(y1-y2)

4×4=16パターンについて、所要時間を計算します。ansの初期値には普通道で $(S_x, S_y)$ から $(G_x, G_y)$ へ移動したときの値を入れておきます。

def solve():
    B,K,Sx,Sy,Gx,Gy = map(int,input().split())
    ans = (abs(Sx-Gx) + abs(Sy-Gy)) * K

    for x1,y1,distS in highway_pos(B,Sx,Sy):
        for x2,y2,distG in highway_pos(B,Gx,Gy):
            ans = min(ans, (distS + distG) * K + highway_dist(B,x1,y1,x2,y2))
    print(ans)

T個のテストケースについて、solve関数を施行します。

T = int(input())
for _ in range(T):
    solve()