日々drdrする人のメモ

今日もdrdr、明日もdrdr

ARC065 - E問題: へんなコンパス / Manhattan Compass

問題: E - へんなコンパス / Manhattan Compass

 {xy}平面上のN個の穴について(a, b)の穴のペアから始めて、マンハッタン距離が同じ穴のペアを辿ること(ペア(p, q)とペア(p, r)の距離が等しければ、一方からもう一方を辿ることができる)で辿れるペア数を求める問題。

今回の解法は、尺取り法とUnion-Findを使う。


まず、マンハッタン距離のひし形が扱いづらいので平面を45度回転させて {\sqrt{2}}倍に拡大する。つまり、 {(x, y) \rightarrow (x-y, x+y)}の変換を行う。
これで、マンハッタン距離 {d(p, q)}の計算を {\max(|x_p - x_q|, |y_p - y_q|)}と見なして計算できる。

この問題では {d(a, b)}と距離が等しい穴のペアのみ辿れるため、今回は距離が {d(a, b)}と同じ穴のペアを平面を走査しながら尺取りで列挙し、 {(a, b)}のペアから辿れる穴をUnion-Findを用いて求め、辿れる穴のペア数を求める。


平面走査はxについてソートした場合とyについてソートした場合それぞれについて行う。
xについてソートした点の列 {P_1, P_2, \ldots, P_N}の平面走査では、ある点 {P_v = (x_v, y_v)}について、 {x_w = x_v - d(a, b)},  {y_v - d(a, b) \le y_w \le y_v + d(a, b)}となる点 {P_w = (x_w, y_w)}を尺取りで全て列挙しながら走査する。この時列挙される点は連続した点の列 {P_s, P_{s+1}, \ldots, P_{t-1}, P_t (s \le t)}となる。

そして列挙した複数の点 {P_w}は、点 {P_v}から辿れる点としてUnion-Findでuniteして管理する。
この時注意すべき点として、ある点 {P_v}で列挙した複数の点 {P_w}をuniteする際、全ての {P_w}について {P_v}とuniteすると、 {O(N^2)}回uniteすることになりTLEする。既にuniteした点同士はuniteする必要がないため、どの点までをuniteし終わったかを持ちながら、点 {P_v}と点 {P_s}のuniteと点 {P_j}と点 {P_{j+1}}  {(s \le j < t)}をuniteすることで {O(N)}回のuniteに抑える。


この平面走査で列挙されるペアには、 {(x_w, y_w) = (x_v - d(a, b), y_v - d(a, b) )}もしくは {(x_w, y_w) = (x_v \pm d(a, b), y_v \mp d(a, b))}となる点のペア {(P_v, P_w)}が2回含まれるため、この重複を除去する必要がある。

最後に、点 {a,b}と同じグループに属する点を含むペアの数を数える。

計算量は {O(N \log{N})}

提出コード(PyPy3): Submission #2114223 - AtCoder Regular Contest 065

N, a, b = map(int, input().split()); a -= 1; b -= 1
P = [] # xでソートする点列
Q = [] # yでソートする点列
for i in range(N):
    x, y = map(int, input().split())
    P.append((x-y, x+y, i))
    Q.append((x+y, x-y, i))

# d(a, b)と見なす距離d
d = max(abs(P[a][0] - P[b][0]), abs(P[a][1] - P[b][1]))

# Union-Find
*parent, = range(N)
def root(x):
    if x == parent[x]:
        return x
    y = parent[x] = root(parent[x])
    return y
def unite(x, y):
    px = root(x); py = root(y)
    if px < py:
        parent[py] = px
    else:
        parent[px] = py
C = [0]*N # C[i]: 点iから重複しないペア数
D = [0]*N # D[i]: 点iから2回重複するペア数

# 重複する点のペアか
def check(P0, i0, j0):
    return abs(P0[i0][0] - P0[j0][0]) == abs(P0[i0][1] - P0[j0][1])

def solve(P0):
    P = P0[:]
    P.sort()

    # prev: どの点までuniteしたか
    s = t = 0; prev = -1
    for i in range(N):
        x, y, i0 = P[i]
        while t < N and P[t][0] < x-d or (P[t][0] == x-d and P[t][1] <= y+d): t += 1
        while s < N and (P[s][0] < x-d or (P[s][0] == x-d and P[s][1] < y-d)): s += 1
        if s < t:
            # P_iと距離dの関係にある点列 P_s, P_{s+1}, ..., P_{t-1}
            j0 = P[s][2]
            unite(i0, j0) # P_i と P_s のunite
            if check(P0, i0, j0):
                D[i0] += 1
            else:
                C[i0] += 1
            if s < t-1:
                j0 = P[t-1][2]
                if check(P0, i0, j0):
                    D[i0] += 1
                    C[i0] += t-s-2
                else:
                    C[i0] += t-s-1
            # P_j と P_{j+1} のunite
            # 既にuniteし終わった点、もしくは点P_sからunite
            for j in range(max(prev, s), t-1):
                unite(P[j][2], P[j+1][2])
            prev = t-1
solve(P)
solve(Q)

# S: 重複しないペア数, T: 2回重複するペア数
S = T = 0
r = root(a)
for i in range(N):
    # 点a,bと同じグループに属するペア数を数える
    if root(i) == r:
        S += C[i]; T += D[i]
print(S + T//2)