[1833] 고속도로 설계하기

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

1833번: 고속철도 설계하기

MST 알고리즘의 기본문제다. 크루스칼 알고리즘을 통해 해결했다.

유니온 함수의 구현에서 사소한 실수를 해서 WA를 받아 당황했지만 실수를 금방 발견했다.

 

풀이

주어지는 간선은 양방향 간선이므로 u (0 ..< N), v (u+1 ..< N) 탐색을 통해 중복되는 정보는 배제하고 탐색하였다.

 

간선이 음수면 이미 연결된 고속도로 이므로 절댓값을 totalCost에 더한 뒤, 두 정점을 연결한다. 주의할 점은 이미 부모가 같은 정점끼리의 간선이 음수로 주어질 수도 있으므로 있으므로 예외처리를 해주어야 한다.

 

나머지의 경우는 edges 배열에 담아낸 뒤 비용이 적은 간선부터 유니온 함수를 수행한다. 이때 연결된 간선은 배열에 따로 정보를 담아주었다가 출력해 주면 된다.

 

정답 코드

	import Foundation
	let N = Int(readLine()!)!
	var arr = Array(repeating: -1, count: N)
	var edges = [(u:Int, v:Int, cost:Int)]()
	var totalCost = 0
	var ans = [(u:Int, v:Int)]()
	var comp = N
	func root(of node:Int) -> Int{
	if arr[node] < 0{
	return node
	}
	arr[node] = root(of: arr[node])
	return arr[node]
	}
	func union(a:Int, b:Int, cost:Int){
	let A = root(of: a)
	let B = root(of: b)
	if A==B { return }
	arr[B] = A
	comp -= 1
	totalCost += cost
	ans.append((a,b))
	}
	for i in 0..<N{
	let cost = readLine()!.split(separator: " ").map{Int($0)!}
	for k in i+1..<N{
	if cost[k] < 0{
	totalCost -= cost[k]
	let u = root(of: i)
	let v = root(of: k)
	if u != v{
	arr[u] = v
	comp -= 1
	}
	}else{
	edges.append((i,k,cost[k]))
	}
	}
	}
	edges.sort(by: {$0.cost > $1.cost})
	while comp>1 && !edges.isEmpty{
	let edge = edges.removeLast()
	union(a: edge.u, b: edge.v, cost: edge.cost)
	}
	print(totalCost, ans.count)
	for ans in ans{
	print(ans.u+1, ans.v+1)
	}

view raw 1833.swift hosted with ❤ by GitHub