https://www.acmicpc.net/problem/23040
23040번: 누텔라 트리 (Easy)
첫째 줄에 트리의 정점의 개수 $N$이 주어진다. ($2 \le N \le 100\,000$) 이후 $(N-1)$개 줄에 걸쳐 각 간선이 잇는 두 정점의 번호 $u_i$, $v_i$가 공백을 사이에 두고 주어진다. ($1 \le u_i \le N$, $1 \le v_i \le N$,
www.acmicpc.net
유니온파인드와 그래프탐색으로 접근했다.
풀이
모든 정점을 방문하게 된다면 시간초과가 일어난다.
예제를 보면 3~1, 3~5, 3~4, 3~2, 6~4, 6~2 총 6개의 구간이 정답이다.
검은색 정점을 시작으로 빨간색 정점으로만 이어지는 경로의 개수. 즉 우리는 연속되는 빨간색 정점을 컴포넌트로 묶어 개수를 미리 저장해 놓으면 보다 빠르게 탐색이 가능해진다.
따라서 인접한 빨간색 정점끼리는 유니온을 통해 하나의 컴포넌트로 묶어내어 개수를 저장하고, 컴포넌트에 인접한 검은색 정점의 개수를 파악하면 해당 컴포넌트에서 발생되는 누텔라 트리의 개수를 구할 수 있게 된다.
해당 과정은 bfs를 통해 O(N)만에 정답을 구할 수 있다.
정답 코드
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import Foundation | |
| // by rhyno | |
| final class FileIO { | |
| private let buffer:[UInt8] | |
| private var index: Int = 0 | |
| init(fileHandle: FileHandle = FileHandle.standardInput) { | |
| buffer = Array(try! fileHandle.readToEnd()!)+[UInt8(0)] // 인덱스 범위 넘어가는 것 방지 | |
| } | |
| @inline(__always) private func read() -> UInt8 { | |
| defer { index += 1 } | |
| return buffer[index] | |
| } | |
| @inline(__always) func readInt() -> Int { | |
| var sum = 0 | |
| var now = read() | |
| var isPositive = true | |
| while now == 10 | |
| || now == 32 { now = read() } // 공백과 줄바꿈 무시 | |
| if now == 45 { isPositive.toggle(); now = read() } // 음수 처리 | |
| while now >= 48, now <= 57 { | |
| sum = sum * 10 + Int(now-48) | |
| now = read() | |
| } | |
| return sum * (isPositive ? 1:-1) | |
| } | |
| @inline(__always) func readString() -> String { | |
| var now = read() | |
| while now == 10 || now == 32 { now = read() } // 공백과 줄바꿈 무시 | |
| let beginIndex = index-1 | |
| while now != 10, | |
| now != 32, | |
| now != 0 { now = read() } | |
| return String(bytes: Array(buffer[beginIndex..<(index-1)]), encoding: .ascii)! | |
| } | |
| } | |
| let file = FileIO() | |
| let N = file.readInt() | |
| var ans = 0 | |
| var arr = Array(repeating: -1, count: N) | |
| var map = Array(repeating: [Int](), count: N) | |
| var visited = Array(repeating: false, count: 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) { | |
| let A = root(of: a) | |
| let B = root(of: b) | |
| if A==B {return} | |
| arr[A] += arr[B] | |
| arr[B] = A | |
| } | |
| for _ in 0..<N-1{ | |
| let u = file.readInt()-1 | |
| let v = file.readInt()-1 | |
| map[u].append(v) | |
| map[v].append(u) | |
| } | |
| var colorOf = file.readString().map{String($0)} | |
| func bfs(from vertex:Int) -> Int{ | |
| var q = [vertex] | |
| var cnt = 0 | |
| var idx = 0 | |
| while idx < q.count{ | |
| let curr = q[idx] | |
| for next in map[curr]{ | |
| if colorOf[next] == "R" && !visited[next]{ | |
| visited[next] = true | |
| union(a: curr, b: next) | |
| q.append(next) | |
| }else if colorOf[next] == "B"{ | |
| cnt += 1 | |
| } | |
| } | |
| idx += 1 | |
| } | |
| return cnt * abs(arr[root(of: vertex)]) | |
| } | |
| for i in 0..<N{ | |
| if colorOf[i] == "R" && !visited[i]{ | |
| visited[i] = true | |
| ans += bfs(from: i) | |
| } | |
| } | |
| print(ans) |
'Problem Solving > BOJ' 카테고리의 다른 글
| [5582] 공통 부분 문자열 (0) | 2024.01.12 |
|---|---|
| [17396] 백도어 (0) | 2023.12.27 |
| [20530] 양분 (0) | 2023.11.04 |
| [15559] 내 선물을 받아줘 (0) | 2023.10.31 |
| [1765] 닭싸움 팀 정하기 (0) | 2023.10.30 |