概要

山岳地帯の地点間を標高差に基づいて水が流れるシミュレーションを行い、ダム設置後の各地点の最終水量を求める問題です。

考察

問題の鍵となる観察は、水が「標高が高い地点から低い地点へしか流れない」という性質です。このため、標高の高い地点から順に処理することで、ある地点を処理する時点ではその地点より高い地点からの水の流入は既に処理済みであることが保証されます。これにより、処理順序を標高の降順に固定することで、効率的なシミュレーションが可能になります。

素朴なアプローチとして、すべての地点を標高順にソートせずに処理すると、水の流れが循環する可能性があり正しい結果が得られません。また、各地点から水を流す際に隣接するすべての地点をチェックすると、計算量が \(O(NM)\) となり、制約条件下では時間内に処理できません。

アルゴリズム

1. 前処理:

   • 各地点について、その地点から直接水が流れることができる「下流隣接地点」（標高が厳密に低い隣接地点）のリストを作成します
   • 各地点の「入次数」（その地点に水が流れ込んでくる経路の数）を計算します（ただし、実際のアルゴリズムでは入次数は使用していません）

2. 処理順序の決定:

   • すべての地点を標高の高い順にソートします。標高が同じ地点間では水が流れないため、これらの順序は結果に影響しません

3. 水量のシミュレーション:

   • ソートされた順序で各地点を処理します
   • 地点にダムがある場合: 何もせず水量を保持します
   • ダムがなく下流隣接地点がある場合: 現在の水量を下流隣接地点数で割り、各下流隣接地点に均等に分配します。自身の水量は0にリセットします
   • ダムがなく下流�接地点がない場合: 水量を保持します

このアプローチにより、高い地点から順に処理することで、各地点を処理する時点ではそれより高い地点からの流入分が既に計算済みであることが保証されます。

計算量

• 時間計算量: \(O(N \log N + M)\)
  • ソートに \(O(N \log N)\)
  • グラフの構築と下流隣接地点の計算に \(O(M)\)
  • 水量のシミュレーションに \(O(N + M)\)
• 空間計算量: \(O(N + M)\)
  • グラフの保存と各種配列のため

実装のポイント

• 入力を効率的に読み取るためにイテレータを使用

• ダム地点は集合で管理して高速な検索を実現

• 浮動小数点数の精度問題に注意（問題文の許容誤差範囲内であれば問題なし）

• 下流隣接地点のみを保持する隣接リストを作成することで、処理を効率化

• 標高順のソート時に、地点のインデックスを保持するためラムダ関数を使用

  ソースコード

import sys
from collections import deque

def main():
    data = sys.stdin.read().split()
    if not data:
        return
    
    it = iter(data)
    N = int(next(it)); M = int(next(it)); K = int(next(it))
    H = [0] * (N+1)
    for i in range(1, N+1):
        H[i] = int(next(it))
    W_initial = [0] * (N+1)
    for i in range(1, N+1):
        W_initial[i] = float(next(it))
    
    graph = [[] for _ in range(N+1)]
    for _ in range(M):
        u = int(next(it)); v = int(next(it))
        graph[u].append(v)
        graph[v].append(u)
    
    dam_set = set()
    if K > 0:
        for _ in range(K):
            s = int(next(it))
            dam_set.add(s)
    
    indegree = [0] * (N+1)
    adj_low = [[] for _ in range(N+1)]
    for u in range(1, N+1):
        for v in graph[u]:
            if H[u] > H[v]:
                adj_low[u].append(v)
                indegree[v] += 1
    
    order = []
    for i in range(1, N+1):
        order.append(i)
    order.sort(key=lambda x: H[x], reverse=True)
    
    water = W_initial[:]
    for u in order:
        if u in dam_set:
            continue
        d = len(adj_low[u])
        if d == 0:
            continue
        flow = water[u] / d
        for v in adj_low[u]:
            water[v] += flow
        water[u] = 0.0
    
    result = []
    for i in range(1, N+1):
        result.append(str(water[i]))
    print(" ".join(result))

if __name__ == "__main__":
    main()

この解説は deepseekv3 によって生成されました。