AtCoder Beginner Contest 245 G問題メモ
G - Foreign Friends
問題
$N$ 頂点 $M$ 辺の重み付き無向グラフが与えられる
頂点 $i$ には値 $A_i$ が書かれている
$N$ 頂点中 $L$ 頂点 $B_{1},B_{2},...,B_{L}$ は「特別な頂点」である
各頂点について、「自分と異なる値が書かれている特別な頂点のうち、最も近いもの」への最短距離を求めよ
$2 \le N \le 10^5$
$1 \le M \le 10^5$
解法
Dijkstraを基本として、2種類のカスタマイズを行う。
複数始点Dijkstra
普通、Dijkstraといえば「単一始点」の最短経路探索アルゴリズムだが、
という場合、最初にキューに突っ込む始点を複数個にしても問題ない。
この問題では、$L$ 個の特別な頂点を始点とし、全頂点への最短距離を求めればよい。
上位k番目までを求める
単なる複数始点Dijkstraで最短距離を求めても、実は始点と終点の値が同じだった！ということがあり得る。
そこで「始点に書かれた値」でグループ化して、上位2グループまでの最短距離を求めれば、
どちらかのグループは終点の値と異なるので、正しい答えとして採用できる。
●特別な頂点、○普通の頂点
❶---1---③---1---②---1---①---4---❸
❷---3---'
❸---5---'
： ： ：
1st (1,❶) (2,❶) (3,❶) (最短距離,始点グループ)
2nd (3,❷) (4,❷) (4,❸)
答え 1 2 4
1stが自身と同じ値の始点グループだったら、2ndを採用すればよい。
このようにしても、計算量は、上位 $k$ グループまでを保持するとして（今回は $k=2$）、通常のDijkstraの $k$ 倍という定数倍で収まる。
ただし、通常のDijkstraで行うような枝刈りは、このカスタマイズにおいても行う必要がある。
実装について
上位 $k$ 番目まで保持するDijkstraにおいては、以下のサイトで言及されているDijkstraの2種類の実装のうち、
1番目を元にした方が書きやすいかも。
というのも、暫定距離配列 dist をいつ更新するかにおいて、
1番目はキューから取り出されたとき
2番目は自身がキューに入れられるとき
なので、2番目は後から更新されることがあり得る。（1番目は確定してからdistに入れるので「暫定」でもない）
後から更新されうると、新しい値で dist をどのように更新すべきかについて
とせねばならず、常に「上位 $k$ 以内の最も劣る値」を取得できなければならない。
ソートなどに計算量を費やすことになり、大変である。
一方、1番目の実装であれば、$k$ 回まで取り出された順に確定していけばよいので、実装も計算量も楽になる。
最短距離のみのDijkstraなら、2番目の実装の方が、キューに突っ込まれる要素を抑えられ、
チェック済みかどうかを管理する配列も持たなくていいのでメモリや計算量的に優しいが、、、
上位 $k$ 番目までを求める場合はそうとは限らないんだね。
Python3
import os
import sys
from heapq import heappop, heappush
import numpy as np
def solve(inp):
n, m, k, l = inp[:4]
aaa = inp[4:4 + n]
bbb = inp[4 + n:4 + n + l] - 1
uuu = inp[4 + n + l::3] - 1
vvv = inp[5 + n + l::3] - 1
ccc = inp[6 + n + l::3]
tuple_list = [(n, n)]
tuple_list.clear()
links = [tuple_list.copy() for _ in range(n)]
for i in range(m):
u = uuu[i]
v = vvv[i]
c = ccc[i]
links[u].append((v, c))
links[v].append((u, c))
def update(dist, nc, na):
if dist[0][1] == na:
if dist[0][0] > nc:
dist[0][0] = nc
return True
elif dist[1][1] == na:
if dist[1][0] > nc:
if dist[0][0] <= nc:
dist[1][0] = nc
else:
dist[1][0] = dist[0][0]
dist[1][1] = dist[0][1]
dist[0][0] = nc
dist[0][1] = na
return True
elif dist[1][0] > nc:
if dist[0][0] <= nc:
dist[1][0] = nc
dist[1][1] = na
else:
dist[1][0] = dist[0][0]
dist[1][1] = dist[0][1]
dist[0][0] = nc
dist[0][1] = na
return True
return False
def check_necessity(distances, cost, u, a):
if distances[u, 0, 0] == cost and distances[u, 0, 1] == a:
return True
if distances[u, 1, 0] == cost and distances[u, 1, 1] == a:
return True
return False
INF = 1 << 60
distances = np.full((n, 2, 2), INF, np.int64)
q = [(n, n, n)]
q.clear()
for b in bbb:
a = aaa[b]
update(distances[b], 0, a)
q.append((0, b, a))
while q:
# print(q)
cost, u, a = heappop(q)
if not check_necessity(distances, cost, u, a):
continue
# print(cost, u, a)
for v, c in links[u]:
nc = cost + c
if update(distances[v], nc, a):
heappush(q, (nc, v, a))
ans = []
for i in range(n):
a = aaa[i]
if distances[i, 0, 1] == INF:
ans.append(-1)
continue
if distances[i, 0, 1] != a:
ans.append(distances[i, 0, 0])
continue
if distances[i, 1, 1] == INF:
ans.append(-1)
continue
ans.append(distances[i, 1, 0])
return ans
SIGNATURE = '(i8[:],)'
if sys.argv[-1] == 'ONLINE_JUDGE':
from numba.pycc import CC
cc = CC('my_module')
cc.export('solve', SIGNATURE)(solve)
cc.compile()
exit()
if os.name == 'posix':
# noinspection PyUnresolvedReferences
from my_module import solve
else:
from numba import njit
solve = njit(SIGNATURE, cache=True)(solve)
print('compiled', file=sys.stderr)
inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(*ans)
