ABC373 D問題の解説（Python）

はじめに

この記事では、ABC373のD問題を解説していきます。

ABC（AtCoder Beginner Contest）とは、AtCoderが開催している、競技プログラミングコンテストです。

ABC373　D – Hidden Weights

問題

N頂点M辺の有向グラフが与えられます。\(j\)番目の有向辺は頂点\(u_j\)から頂点\(v_j\)に向かっており、重み\(w_j\)を持っています。

各頂点に整数を書き込む方法であって、次の条件を満たすものを1つ見つける問題です。

頂点\(i\)に書き込まれている値を\(x_i\)とする。すべての辺 \(j=1,2,…,M\) について、\(x_{v_j}-x_{u_j}=w_j\)が成り立つ。

なお、与えられる入力について、条件を満たす書き込み方が少なくとも1つ存在することが保証されているものとします。

制約

\(2 \leqq N \leqq 2×10^5 \)
\(1 \leqq M \leqq min\{2×10^5, N(N-1)/2\} \)
\(1 \leqq u_i, v_j \leqq N\)
\(u_i≠v_j\)
\(i≠j\) なら \((u_i, v_i)≠(u_j, v_j)\) かつ \((u_i, v_i)≠(v_j, u_j)\)
\(|w_j| \leqq 10^9 \)

思考の筋道

グラフのそれぞれの連結成分において、いずれかの頂点の値を決めれば、残りの頂点の値も決まります。

ここでは、DFSを用いた方法と、UnionFindを用いた方法を説明します。

DFS（深さ優先探索）を使用する場合

グラフの隣接リストを作る際、隣接する頂点とセットで重みの情報も渡しましょう。そうすることで、深さ優先探索を進める中で、頂点の値を決めていくことができます。

DFSはスタックを用いて実装しています。スタックに最後に入ってきたものから取り出すことで、深さ優先探索を実現しているということですね。また、調査済みの頂点を管理しておき、同じ頂点を2度調べないようにしています。

一般的なDFSのコードは手元に持っておくとよいでしょう。

UnionFind を使用する場合

連結成分に関する問題なので、UnionFindを思い浮かべた方もいるのではないかと思います。実は、重み付きUnionFindというものを使うことによって、この問題を解くことができます。

通常のUnionFindとの違いは、各頂点について、根との重みの差 diff_weight の情報を持っていることです。

自身の根を求めるroot関数内で、diff_weightも計算できるようにします。

また、unite関数で二つの連結成分を併合する際に、もともと根であった頂点が根でなくなることがあります。その際に diff_weight の更新が必要ですね。

UnionFindのコードも手元に持っておくとよいでしょう。

コード

DFS（深さ優先探索）を使用する場合

def dfs(g, u):
    if not visited[u]:
        stack = [u]
        visited[u] = True
        while stack:
            u = stack.pop()
            for v, w in g[u]:
                if not visited[v]:
                    ans[v] = ans[u] + w
                    stack.append(v)
                    visited[v] = True
    return visited

N, M = map(int, input().split())
l = [list(map(int, input().split())) for _ in range(M)]

G = [[] for _ in range(N)]
for row in l:
    u, v, w = row[0] - 1, row[1] - 1, row[2]
    G[u].append((v, w))
    G[v].append((u, -w))

ans = [0] * N
visited = [False] * N

for i in range(N):
    dfs(G, i)
print(*ans)

（重み付き）UnionFind を使用する場合

from collections import defaultdict

class WeightedUnionFind():
    def __init__(self, n):
        self.n = n
        self.par = [-1] * (n + 1) 
        self.rank = [0] * (n + 1)
        self.diff_weight = [0] * (n + 1)

    # 根を求める
    def root(self, x):
        if self.par[x] == -1:
            return x
        y = self.root(self.par[x])
        self.diff_weight[x] += self.diff_weight[self.par[x]]
        self.par[x] = y
        return y

    # 重みを取得する
    def weight(self, x):
        self.root(x)
        return self.diff_weight[x]

    # xとyが同じグループに属するか
    def issame(self, x, y):
        return self.root(x) == self.root(y)

     # xを含む連結成分とyを含む連結成分を併合する
    def unite(self, x, y, w):
        w += self.weight(x) - self.weight(y)
        rx = self.root(x)
        ry = self.root(y)
        if rx == ry:
            return
        if self.rank[rx] < self.rank[ry]:
            rx, ry = ry, rx
            w = -w
        if self.rank[rx] == self.rank[ry]:
            self.rank[rx] += 1
        self.par[ry] = rx
        self.diff_weight[ry] = w

    # xとyの間の重み差を計算する
    def diff(self, x, y):
        if not self.issame(x, y):
            return None
        return self.weight(y) - self.weight(x)

N, M = map(int, input().split())
l = [list(map(int, input().split())) for _ in range(M)]

uf = WeightedUnionFind(N)

for row in l:
    uf.unite(row[0] - 1, row[1] - 1, row[2])

ans = [0] * N

for i in range(N):
    root = uf.root(i)
    ans[i] = uf.diff(root, i)

print(*ans)

以上、ABC373のD問題の解説でした！

では、またね。

リンク

ABC373

情報科学の部屋へ

ABC373 D問題の解説（Python）

はじめに