AtCoder Beginner Contest 245 G問題メモ

AtCoder Beginner Contest 245

今年度も終わりだね

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$ 以内でない場合、スキップし、無駄に隣接頂点を調べない
キューに入れようとする時点で既に上位 $k$ 以内でない場合、スキップし、無駄な値を入れない

実装について

上位 $k$ 番目まで保持するDijkstraにおいては、以下のサイトで言及されているDijkstraの2種類の実装のうち、 1番目を元にした方が書きやすいかも。

二種類のDijkstra法とグラフの負辺 - 私と理論

というのも、暫定距離配列 dist をいつ更新するかにおいて、

1番目はキューから取り出されたとき
2番目は自身がキューに入れられるとき

なので、2番目は後から更新されることがあり得る。（1番目は確定してからdistに入れるので「暫定」でもない）

後から更新されうると、新しい値で dist をどのように更新すべきかについて

上位 $k$ 以内に既に自グループがある場合、その値との比較で小さければ更新
上位 $k$ 以内に自グループがない場合、その中で最も劣る値との比較で小さければそいつを追い出して更新

とせねばならず、常に「上位 $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)

            

        
      