トップページ -> AOJの解答例 -> GRL_5の解答例

GRL_5の解答例(Python)

木に関する問題です．

1. GRL_5_A: Diameter of a Tree

最遠頂点間の距離(木の直径)を求める問題です．適当な点から最も遠い点を求めた後に，その点から最も遠い点との距離を求めます．最も遠い点は深さ優先探索で求めることができます．


# GRL_5_A
# 深さ優先探索で一番深い2点間の距離を調べればよい
from collections import deque

V = int(input())

# グラフを隣接リストとして受け取る
adjacent_list = [[] for _ in range(V)]
for _ in range(V-1):
    s,t,w = map(int,input().split())
    adjacent_list[s].append((w,s,t)) 
    adjacent_list[t].append((w,t,s))

# 深さ優先探索
def DFS(u):
    # print(u)
    # uからの距離の辞書
    dist_dict = {i:-1 for i in range(V)}
    dist_dict[u] = 0
    # 巡回済みの点の集合
    done_set = set([u])
    # 走査すべき辺のキュー
    queue = deque(adjacent_list[u])
    # 現在分かっている最大の距離と節点
    max_dist = 0
    max_point = 0
    while queue:
        w,s,t = queue.pop()
        
        # u からの距離を更新する
        dist_dict[t] = max(dist_dict[s]+w,dist_dict[t])
        # 距離が現在分かっている最大より大きければ更新
        if max_dist < dist_dict[t]:
            max_point = t
            max_dist = dist_dict[t]
        if t not in done_set:
            done_set.add(t)
            for info in adjacent_list[t]:
                _,_,t2 = info
                if t2 not in done_set:
                    queue.append(info)
    
    return max_dist,max_point
    
# 0から最も遠い点を探す
_,max_point = DFS(0)
# 0から最も遠い点から最も遠い点を探す
max_dist,_ = DFS(max_point)
print(max_dist)

2. GRL_5_B: Height of a Tree

各点から一番遠い点との距離を求める問題です． DFS(u)をすべての点について走らせれば答えは出るのですが，間に合わないので少し工夫します．各点から一番遠い点は木の直系を求めるときに用いた2点のどちらかなのでDFS(u)をすこし弄ってuからの距離をdist_dictで返すようにしました． 2つのdist_dictを比べて距離の大きいほうが最も遠い点からの距離です．


# GRL_5_B
# 深さ優先探索で一番深い2点間の距離を調べればよい
from collections import deque

V = int(input())

# グラフを隣接リストとして受け取る
adjacent_list = [[] for _ in range(V)]
for _ in range(V-1):
    s,t,w = map(int,input().split())
    adjacent_list[s].append((w,s,t)) 
    adjacent_list[t].append((w,t,s))

# 深さ優先探索
def DFS(u):
    # print(u)
    # uからの距離の辞書
    dist_dict = {i:-1 for i in range(V)}
    dist_dict[u] = 0
    # 巡回済みの点の集合
    done_set = set([u])
    # 走査すべき辺のキュー
    queue = deque(adjacent_list[u])
    # 現在分かっている最大の距離と節点
    max_dist = 0
    max_point = 0
    while queue:
        w,s,t = queue.pop()
        
        # u からの距離を更新する
        dist_dict[t] = max(dist_dict[s]+w,dist_dict[t])
        # 距離が現在分かっている最大より大きければ更新
        if max_dist < dist_dict[t]:
            max_point = t
            max_dist = dist_dict[t]
        if t not in done_set:
            done_set.add(t)
            for info in adjacent_list[t]:
                _,_,t2 = info
                if t2 not in done_set:
                    queue.append(info)
    
    return max_dist,max_point,dist_dict
    
# 0から最も遠い点を探す
max_dist,max_point,dist_dict = DFS(0)

# 最遠点の片方からの距離
max_dist,max_point,dist_dict1 = DFS(max_point)

# 最遠点のもう片方からの距離
max_dist,max_point,dist_dict2 = DFS(max_point)

# どっちかが必ず一番遠い点になる
for i in range(V):
    print(max(dist_dict1[i],dist_dict2[i]))

3. GRL_5_C: Lowest Common Ancestor

根から最も遠い共通の先祖(最小共通先祖)を求める問題です． UnionFindで互いに素な集合を管理しながら求めます．
※ アルゴリズムイントロダクション第22章(P158)


# GRL_5_C
# DSL_1_A の Union-Findを流用
class UnionFind:
    def __init__(self, n):
        # 親を格納するリストと各要素のランクを格納するリストを初期化する
        self.parent = [i for i in range(n)]  # 最初は自分自身が親
        self.rank = [0] * n  # ランクは最初すべて0
    
    def find(self, x):
        # 要素xの親を見つける
        if self.parent[x] != x:
            # 要素の親が自分自身でない場合、再帰的に親を探す
            self.parent[x] = self.find(self.parent[x])  # 親を再帰的に探して更新
        return self.parent[x]
    
    def unite(self, x, y):
        # xとyのルートを見つける
        root_x = self.find(x)
        root_y = self.find(y)
        
        if root_x == root_y:
            # xとyが同じグループに属している場合は何もしない
            return None
        
        # ランクを考慮して結合する
        if self.rank[root_x] < self.rank[root_y]:
            # root_xのランクがroot_yのランクより小さい場合、root_xをroot_yの子にする
            self.parent[root_x] = root_y
        elif self.rank[root_x] > self.rank[root_y]:
            # root_xのランクがroot_yのランクより大きい場合、root_yをroot_xの子にする
            self.parent[root_y] = root_x
        else:
            # 両方のランクが同じ場合はどちらを親にしても良いが、ここではroot_yをroot_xの子にする
            self.parent[root_y] = root_x
            self.rank[root_x] += 1  # root_xのランクを増やす
            
    def isSame(self,x,y):
        return self.find(x) == self.find(y)
            
# アルゴリズムイントロダクション 第22章(P158)
ancestor = {}
done_set = set()
ans_dict = {}
def LCA(u):
    if u in done_set:
        return False
    done_set.add(u)
    
    ancestor[UF.find(u)] = u
    for info in adjacent_list[u]:
        v = info[2]
        LCA(v)
        UF.unite(u,v)
        ancestor[UF.find(u)] = u
        
    done_set.add(u)
    
    if u in q_dict:
        for v in q_dict[u]:
            if v in done_set:
                ans_dict[f"{u} {v}"] = ancestor[UF.find(v)]
                ans_dict[f"{v} {u}"] = ancestor[UF.find(v)]
                
import sys
sys.setrecursionlimit(10**8) # 再帰回数の変更(LCAは再帰関数なので)
            
V = int(input())
UF = UnionFind(V) # union-find 詳しくはDSL1

# グラフを隣接リストとして受け取る
adjacent_list = [[] for _ in range(V)]
for i in range(V):
    info = list(map(int,input().split()))
    for t in info[1:]:
        adjacent_list[i].append((1,i,t)) # 重みは常に1にしておく

# クエリの辞書とリストを作っておく 
# リストは順序を保持したいだけなので stdin.readline を使ってもいい
q_dict = {}
q_list = []
Q = int(input())
for _ in range(Q):
    u,v = map(int,input().split())
    if u not in q_dict:
        q_dict[u] = set()
    q_dict[u].add(v)
    
    if v not in q_dict:
        q_dict[v] = set()
    q_dict[v].add(u)
    
    q_list.append((u,v))

# 最小共通先祖を求める
LCA(0)

# クエリの順序に従って結果を出力
for q in q_list:
    u,v = q
    print(ans_dict[f"{u} {v}"])

<- 前へ戻る【目次に戻る】次へ進む ->