ARC052 - C問題: 高橋くんと不思議な道 - 日々ｄｒｄｒする人のメモ

N個の頂点と２つのタイプを含むM個の辺が存在する。タイプAは1のコスト、タイプBは(今まで通過したタイプBの辺数)+1のコストがかかる時、頂点0から各頂点に到達するための最小コストはいくらかを答える問題。

コンテスト中にも解けなかった問題。久しぶりにやっても解けず、想定解法(+α)で通したメモ。

想定解法の解説は公式の解説を参照。

今回のダイクストラは、(頂点i, 何回タイプBを通過したか)の状態を1つの頂点として探索を行うが、そのままだと状態数が ${O(N^2)}$ でTLEする。

想定解法は、各頂点iに到達するために必要な最小のタイプBの辺数 ${b_i}$ を計算しておき、その上で各頂点に到達する時に通過してきたタイプBの辺数aが ${\frac{a(a+1)}{2} \le \frac{b_i(b_i+1)}{2} + N - 1}$ を満たすようにダイクストラ法を行うと状態数が ${O(N\sqrt{N})}$ 個程度になり、通るようになるというもの。

しかし、この想定解法では、PyPy3, Python3ではTLEする。
提出: Submission #2119297 - AtCoder Regular Contest 052

想定解法は、各頂点iのコスト上限を ${\frac{b_i(b_i+1)}{2} + N - 1}$ としたものとほぼ同じであるが、この上限は大きいためTLEする(タイプBの辺が多く、 ${b_i }$ が小さい頂点が多いケースが重いはず)。そこで、このコスト上限を少し強めに見積もることを考える。
今回は、「出来る限りタイプBの辺を使わずに」各頂点に到達するための最小コストを上限とする。
このコストは通常のダイクストラで計算できる。(例えば、タイプBのコストをNとして計算し、得られたコスト ${c_i}$ から ${(c_i \% N) + \frac{\lfloor c_i / N \rfloor (\lfloor c_i / N \rfloor + 1)}{2}}$ を計算する等)

この上限は想定解法の上限よりも小さくなる。ある頂点に(M - p)個のタイプAの辺とp個のタイプBの辺 ${(M, p \le N - 1)}$ を利用することで最小コストで到達できる時、コスト上限は ${\frac{p(p+1)}{2} + (M - p)}$ となる。こちらも状態数は ${O(N\sqrt{N})}$ 個だが定数は小さくなる。

この解法によりPython3で通すことができる。
提出: Submission #2119110 - AtCoder Regular Contest 052

N, M = map(int, input().split())
G = [[] for i in range(N)]
for i in range(M):
    c, a, b = map(int, input().split())
    G[a].append((b, c))
    G[b].append((a, c))

from heapq import heappush, heappop

INF = 10**18

que = [(0, 0)]
cost_m = [INF]*N # cost_m[i]: 各頂点iにおけるコスト上限
b_min = [INF]*N
cost_m[0] = b_min[0] = 0

# 各頂点におけるコスト上限を求めるダイクストラ
while que:
    cost, v = heappop(que)
    if cost_m[v] < cost:
        continue

    for w, c in G[v]:
        if c == 0:
            # タイプA: コスト1
            if cost + 1 < cost_m[w]:
                cost_m[w] = cost + 1
                heappush(que, (cost + 1, w))
        else:
            # タイプB: コストNとみなす
            if cost + N < cost_m[w]:
                cost_m[w] = cost + N
                heappush(que, (cost + N, w))

# 結果から各頂点の上限コストを計算
for i in range(N):
    b_min[i] = b = cost_m[i] // N
    cost_m[i] = ((cost_m[i] % N) + b*(b+1)//2)+1

# コスト上限以内で解を求めるダイクストラ
que = [(0, 0, 0)]
# dist[i][k]: 頂点iで、タイプB辺をk回通過した時の最小コスト
dist = [{} for i in range(N)]
dist[0][0] = 0
while que:
    cost, v, k = heappop(que) # (コスト, 頂点v, タイプBの辺の通過回数)
    if dist[v][k] < cost:
        continue
    for w, c in G[v]:
        if c == 0:
            # タイプA: dist[v][k] --> dist[w][k]
            if cost + 1 < dist[w].get(k, cost_m[w]):
                dist[w][k] = cost + 1
                heappush(que, (cost + 1, w, k))
        else:
            # タイプB: dist[v][k] --> dist[w][k+1]
            if cost + k + 1 < dist[w].get(k+1, cost_m[w]):
                dist[w][k+1] = cost + k + 1
                heappush(que, (cost + k + 1, w, k+1))
print(*(min(d.values()) for d in dist), sep='\n')

コンテスト中に探索量を減らすための考察力が欲しい。