[백준/Python] 32690번 Starlight Express

문제

https://www.acmicpc.net/problem/32690

 

소스코드

import sys
input = sys.stdin.readline
n = int(input())
parent = [i for i in range(n + 1)]
rank = [0] * (n + 1)
size = [1] * (n + 1)  # 각 집합의 크기
xy = [[0, 0] for _ in range(n + 1)]
dict_r_x = {}
dict_r_y = {}
bu = set()

def find(x):
    if x != parent[x]:
        parent[x] = find(parent[x])
    return parent[x]

def union(x, y):
    x_root = find(x)
    y_root = find(y)

    if x_root != y_root:
        if rank[x_root] < rank[y_root]:
            parent[x_root] = y_root
            size[y_root] += size[x_root]
        else:
            parent[y_root] = x_root
            size[x_root] += size[y_root]
            if rank[x_root] == rank[y_root]:
                rank[x_root] +=1

def is_connected(x, y):
    return find(x) == find(y)

for i in range(1, n + 1):
    x, y = map(int, input().split())
    if x not in dict_r_x:
        dict_r_x[x] = []
    if y not in dict_r_y:
        dict_r_y[y] = []
    dict_r_x[x].append(i)
    dict_r_y[y].append(i)
    xy[i][0] = x 
    xy[i][1] = y


for stations in dict_r_x.values():
    first_station = stations[0]
    for station in stations[1:]:
        union(first_station, station)

for stations in dict_r_y.values():
    first_station = stations[0]
    for station in stations[1:]:
        union(first_station, station)

root_sizes = {}
for i in range(1, n+1):
    root = find(i)
    if root not in root_sizes:
        root_sizes[root] = size[root]

roots_by_size = sorted(root_sizes.items(), key=lambda x: -x[1])
most_common = roots_by_size[:2]

if len(most_common) == 1:
    x = xy[most_common[0][0]][0]
    y_candidate = -10**9
    while (x, y_candidate) in bu:
        y_candidate += 1
    print(x, y_candidate)
else:
    x1 = xy[most_common[0][0]][0]
    y2 = xy[most_common[1][0]][1]
    print(x1, y2)

 

유니온파인드를 이용한 풀이입니다. 

 

아래는 유니온 파인드를 알고 있다고 가정한 풀이입니다. 

 

입력으로 n과 n개의 x, y 좌표가 주어집니다. 

역은 같은 x와 y 좌표에 있는 역이라면 이동할 수 있습니다. 

 

이 문제에서는 단 하나의 좌표를 추가해 최대한 많은 역들을 연결해야합니다. 

 

x, y 좌표가 같은 역들을 모두 같은 그룹으로 지정한 뒤 (union) 

만약 모든 역이 같은 그룹이 아니라면 사이즈가 제일 큰 그룹과 두번째로 큰 그룹의 각 대표 역의 x, y 좌표을 출력하면 정답입니다. 

만약 모든 역이 같은 그룹이라면 아무 역의 x나 y가 같은 좌표 중 입력에서 등장하지 않은 좌표를 출력하면 됩니다.    

 

----------------------------------------------- 상세 풀이 ------------------------------------------------------

import sys
input = sys.stdin.readline

 

파이썬의 빠른 입출력을 책임져줄 친구입니다. 

 

parent = [i for i in range(n + 1)]
rank = [0] * (n + 1)
size = [1] * (n + 1)

def find(x):
    if x != parent[x]:
        parent[x] = find(parent[x])
    return parent[x]

def union(x, y):
    x_root = find(x)
    y_root = find(y)

    if x_root != y_root:
        if rank[x_root] < rank[y_root]:
            parent[x_root] = y_root
            size[y_root] += size[x_root]
        else:
            parent[y_root] = x_root
            size[x_root] += size[y_root]
            if rank[x_root] == rank[y_root]:
                rank[x_root] +=1

def is_connected(x, y):
    return find(x) == find(y)

 

랭킹기반 유니온파인드입니다. 문제 특성상 union할때 size를 동시에 계산합니다

 

xy = [[0, 0] for _ in range(n + 1)]
dict_r_x = {}
dict_r_y = {}
bu = set()

for i in range(1, n + 1):
    x, y = map(int, input().split())
    bu.add((x, y))
    if x not in dict_r_x:
        dict_r_x[x] = []
    if y not in dict_r_y:
        dict_r_y[y] = []
    dict_r_x[x].append(i)
    dict_r_y[y].append(i)
    xy[i][0] = x 
    xy[i][1] = y

 

입력받습니다. 

set인 bu는 등장한 모든 역의 좌표를 저장합니다. (출력할때 한번도 등장하지않은 역을 출력해야하기에)

이중 list인 xy는 역의 빠른 좌표 접근을 위해 존재합니다.  - xy[역번호][0은 x, 1은 y]

딕셔너리인 dict_r_x와 dict_r_y는 각 x, y좌표를 키로 사용해 역을 저장합니다. (역을 빠르게 union하기위해)

 

for stations in dict_r_x.values():
    first_station = stations[0]
    for station in stations[1:]:
        union(first_station, station)

for stations in dict_r_y.values():
    first_station = stations[0]
    for station in stations[1:]:
        union(first_station, station)

 

x좌표, y좌표가 같은 모든 역들을 union해줍니다.

 

root_sizes = {}
for i in range(1, n+1):
    root = find(i)
    if root not in root_sizes:
        root_sizes[root] = size[root]
        
roots_by_size = sorted(root_sizes.items(), key=lambda x: -x[1])
most_common = roots_by_size[:2]

 

각 역에 대해 find(i) 로 경로합축을 합니다. (자기의 가장 높은 조상을 똑바로 볼 수 있게)

각 루트의 사이즈를 크기 기준 내림차순으로 정렬합니다. 

 

if len(most_common) == 1:
    x = xy[most_common[0][0]][0]
    y_candidate = -10**9
    while (x, y_candidate) in bu:
        y_candidate += 1
    print(x, y_candidate)
else:
    x1 = xy[most_common[0][0]][0]
    y2 = xy[most_common[1][0]][1]
    print(x1, y2)

 

만약 most_common이 1이라면 ( 즉 모든 역이 모두 연결되어있다면)

그룹의 대표자의 x좌표를 가져오고 y좌표를 -10**9 (문제에서 주어진 가장 작은 좌표)부터 하나씩늘리며 앞서 모든 역의 좌표를 저장한 bu 에 존재 하는지 체크한 후 존재하지않으면 출력합니다. 

 

만약 most_common이 1이 아니라면 

가장 많은 그룹의 대표역의 x좌표와 두번째로 가장 많은 그룹의 y좌표를 조합하여 출력합니다.