https://www.acmicpc.net/problem/17472
삼성 SW 역량테스트 A+을 좌절시킨 그 문제이다. 테스트를 볼 때는 최소 신장 트리 개념조차 몰라서 풀지 못했던 문제..
풀어보도록 한다.
먼저 각 섬마다 번호를 붙여주기 위해서 BFS를 진행합니다. 한 번의 BFS가 끝나면 번호를 1 증가시킵니다.
answer = 0
n, m = map(int, input().split())
islands = [list(input().split()) for _ in range(n)]
check = [[0] * m for _ in range(n)]
dirs = [(0, 1), (1, 0), (0, -1), (-1, 0)]
num = 1
for i in range(n):
for j in range(m):
if islands[i][j] == '1' and not check[i][j]:
check[i][j] = num
q = []
q.append([i, j])
while q:
x, y = q.pop(0)
for dir in dirs:
dx, dy = x + dir[0], y + dir[1]
if 0 <= dx < n and 0 <= dy < m:
if islands[dx][dy] == '1' and not check[dx][dy]:
check[dx][dy] = num
q.append([dx, dy])
num += 1
그 다음 각 섬 사이의 최소 거리, edge의 최소값을 찾기 위해 각 섬의 테두리를 찾습니다. 4방향을 탐색해 0인 부분이 하나라도 있다면 섬의 테두리, 즉 edge를 생성할 수 있다는 뜻입니다.
이후에는 key : (섬1, 섬2) / value : edge 가중치 형태의 딕셔너리(tree_dict)를 이용해 섬과 섬이 이어지는 부분을 찾을때마다 value를 더 작은 값으로 갱신합니다.
tree_dict = dict()
visited = [[False] * m for _ in range(n)]
for i in range(n):
for j in range(m):
if check[i][j] and not visited[i][j]:
bridge_dict = dict()
visited[i][j] = True
q = [[i, j]]
while q:
flag = False
x, y = q.pop(0)
for dir in dirs:
dx, dy = x + dir[0], y + dir[1]
if 0 <= dx < n and 0 <= dy < m:
if not check[dx][dy]:
flag = True
if flag:
for dir in dirs:
dx, dy = x + dir[0], y + dir[1]
if 0 <= dx < n and 0 <= dy < m:
if not visited[dx][dy]:
if not check[dx][dy]:
bridge_dict[(x, y)] = bridge_dict.get((x, y), []) + [[dx-x, dy-y]]
else:
visited[dx][dy] = True
q.append([dx, dy])
visited[x][y] = True
q = []
for key, value in bridge_dict.items():
for v in value:
q.append([(key[0], key[1]), (v[0], v[1]), 0])
while q:
info = q.pop(0)
x, y, dirx, diry, distance = info[0][0], info[0][1], info[1][0], info[1][1], info[2]
dx, dy = x + dirx, y + diry
if 0 <= dx < n and 0 <= dy < m:
if not check[dx][dy]:
q.append([(dx, dy), (dirx, diry), distance+1])
else:
i1, i2 = check[i][j], check[dx][dy]
if i1 != i2 and distance > 1:
tree_dict[(min(i1, i2), max(i1, i2))] = min(tree_dict.get((min(i1, i2), max(i1, i2)), 10), distance)
dictionary를 list 형태로 만든 뒤, 가중치의 오름차순으로 정렬합니다. size를 이용해 MST의 효율성을 높이겠습니다.
각 Node가 다른 부모를 갖고 있는 경우 MST를 만족하기 때문에 size와 parent를 변경한 후 다음 가중치의 Node 2개를 탐색합니다.
num이 6번이라면 섬은 5개가 존재하고, MST의 Edge가 4개가 되면 완성되므로 반복문을 탈출합니다(MST의 edge 갯수는 Node 갯수-1개입니다).
모든 Node를 순회하거나 반복문을 탈출했을 때 edge 갯수가 MST를 충족하지 않거나(count != num-2), MST를 생성하는 edge가 존재하지 않을 경우(answer == 0) -1을 출력하고 그 외에는 answer를 출력합니다.
tree_list = [(key,value) for key, value in tree_dict.items()]
tree_list.sort(key=lambda x : x[1])
size = [1 for _ in range(num+1)]
parent = [x for x in range(num+1)]
count = 0
for node, edge in tree_list:
if count == num-2: break
n1, n2 = node[0], node[1]
p1, p2 = mst(n1, parent), mst(n2, parent)
if p1 != p2:
s1, s2 = size[n1], size[n2]
if s1 > s2:
parent[p1] = p2
size[n2] += size[n1]
else:
parent[p2] = p1
size[n1] += size[n2]
count += 1
answer += edge
if answer == 0 or count != num-2:
answer = -1
전체 코드(Python)
def mst(node, parent):
if node == parent[node]:
return node
root = mst(parent[node], parent)
return root
answer = 0
n, m = map(int, input().split())
islands = [list(input().split()) for _ in range(n)]
check = [[0] * m for _ in range(n)]
dirs = [(0, 1), (1, 0), (0, -1), (-1, 0)]
num = 1
for i in range(n):
for j in range(m):
if islands[i][j] == '1' and not check[i][j]:
check[i][j] = num
q = []
q.append([i, j])
while q:
x, y = q.pop(0)
for dir in dirs:
dx, dy = x + dir[0], y + dir[1]
if 0 <= dx < n and 0 <= dy < m:
if islands[dx][dy] == '1' and not check[dx][dy]:
check[dx][dy] = num
q.append([dx, dy])
num += 1
tree_dict = dict()
visited = [[False] * m for _ in range(n)]
for i in range(n):
for j in range(m):
if check[i][j] and not visited[i][j]:
bridge_dict = dict()
visited[i][j] = True
q = [[i, j]]
while q:
flag = False
x, y = q.pop(0)
for dir in dirs:
dx, dy = x + dir[0], y + dir[1]
if 0 <= dx < n and 0 <= dy < m:
if not check[dx][dy]:
flag = True
if flag:
for dir in dirs:
dx, dy = x + dir[0], y + dir[1]
if 0 <= dx < n and 0 <= dy < m:
if not visited[dx][dy]:
if not check[dx][dy]:
bridge_dict[(x, y)] = bridge_dict.get((x, y), []) + [[dx-x, dy-y]]
else:
visited[dx][dy] = True
q.append([dx, dy])
visited[x][y] = True
q = []
for key, value in bridge_dict.items():
for v in value:
q.append([(key[0], key[1]), (v[0], v[1]), 0])
while q:
info = q.pop(0)
x, y, dirx, diry, distance = info[0][0], info[0][1], info[1][0], info[1][1], info[2]
dx, dy = x + dirx, y + diry
if 0 <= dx < n and 0 <= dy < m:
if not check[dx][dy]:
q.append([(dx, dy), (dirx, diry), distance+1])
else:
i1, i2 = check[i][j], check[dx][dy]
if i1 != i2 and distance > 1:
tree_dict[(min(i1, i2), max(i1, i2))] = min(tree_dict.get((min(i1, i2), max(i1, i2)), 10), distance)
tree_list = [(key,value) for key, value in tree_dict.items()]
tree_list.sort(key=lambda x : x[1])
size = [1 for _ in range(num+1)]
parent = [x for x in range(num+1)]
count = 0
for node, edge in tree_list:
if count == num-2: break
n1, n2 = node[0], node[1]
p1, p2 = mst(n1, parent), mst(n2, parent)
if p1 != p2:
s1, s2 = size[n1], size[n2]
if s1 > s2:
parent[p1] = p2
size[n2] += size[n1]
else:
parent[p2] = p1
size[n1] += size[n2]
count += 1
answer += edge
if answer == 0 or count != num-2:
answer = -1
print(answer)
'알고리즘 > 백준' 카테고리의 다른 글
백준 14503 로봇청소기(Python, C++) (0) | 2020.03.21 |
---|