概要

都市間の最短距離と代表者の忍耐力を用いて、条件を満たす「友好ペア」の総数を求める問題です。

考察

この問題では、まず各都市間の最短距離を求める必要があります。都市数 \(N\) は最大 400 なので、全点対間の最短距離を求めるのに適したアルゴリズムとして Dijkstra法 を各頂点から実行する方法が考えられます（Floyd-Warshall法より高速かつ実装が簡単です）。

次に、各代表者 \(i\) について、他のすべての都市 \(j\) に対して \(d(i, j) > S_i\) であるかどうかを判定し、そのような都市の数をカウントして 不満の個数 \(C_i\) を求めます。

最後に、すべての異なる代表者ペア \((i, j)\)（\(i < j\)）について、以下の2条件を満たすものを「友好ペア」としてカウントします：

1. \(d(i, j) \leq S_i\) かつ \(d(i, j) \leq S_j\)：互いに相手の都市に不満を抱かない。
2. \(C_i = C_j\)：両者の不満の個数が等しい。

素朴な解法では、各頂点からのDijkstraをナイーブに行うと全体で \(O(N \cdot (N + M) \log N)\) となりますが、制約上これでも十分高速です。逆に、Floyd-Warshallを使うと \(O(N^3)\) かかり、TLEする可能性があります。

アルゴリズム

1. グラフの構築
   入力から無向グラフを隣接リスト形式で構築します。

2. 全点対最短距離の計算
   各頂点 \(i\) に対して Dijkstra 法を実行し、\(d(i, j)\) を求めたテーブル D を作成します。

3. 不満の個数の計算
   各代表者 \(i\) について、\(d(i, j) > S_i\) となる都市 \(j\) の数を数え、配列 C に保存します。

4. 友好ペアのカウント
   全ての \(i < j\) について、以下の条件を満たすペアを数えます：

   • \(D[i][j] \leq S[i]\) かつ \(D[i][j] \leq S[j]\)
   • \(C[i] == C[j]\)

計算量

• 時間計算量: \(O(N \cdot (N + M) \log N)\)
• 空間計算量: \(O(N^2 + M)\)

実装のポイント

• 頂点番号は入力で1-indexedで与えられるため、内部処理では0-indexedに変換する必要があります。

• Dijkstra法では優先度付きキュー（heapq）を使用して効率的に最小距離頂点を取り出すようにしましょう。

• 無限大を表す値として float('inf') を使用しており、比較演算も正しく動作します。

  ソースコード

import heapq
import sys

def dijkstra(graph, start, n):
    dist = [float('inf')] * n
    dist[start] = 0
    pq = [(0, start)]
    while pq:
        d, u = heapq.heappop(pq)
        if d > dist[u]:
            continue
        for v, w in graph[u]:
            if dist[u] + w < dist[v]:
                dist[v] = dist[u] + w
                heapq.heappush(pq, (dist[v], v))
    return dist

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    
    idx = 0
    N = int(data[idx]); idx += 1
    M = int(data[idx]); idx += 1
    S = [int(data[idx + i]) for i in range(N)]
    idx += N
    
    graph = [[] for _ in range(N)]
    for _ in range(M):
        u = int(data[idx]) - 1; idx += 1
        v = int(data[idx]) - 1; idx += 1
        w = int(data[idx]); idx += 1
        graph[u].append((v, w))
        graph[v].append((u, w))
    
    # 全点対最短距離を計算
    D = []
    for i in range(N):
        dist = dijkstra(graph, i, N)
        D.append(dist)
    
    # 各代表者の不満の個数 C_i を計算
    C = [0] * N
    for i in range(N):
        count = 0
        for j in range(N):
            if i == j:
                continue
            if D[i][j] > S[i]:
                count += 1
        C[i] = count
    
    # 友好ペアの数を数える
    ans = 0
    for i in range(N):
        for j in range(i + 1, N):
            if D[i][j] <= S[i] and D[i][j] <= S[j] and C[i] == C[j]:
                ans += 1
    print(ans)

if __name__ == "__main__":
    main()

この解説は qwen3-coder-480b によって生成されました。