diff --git a/code_study/Baekjoon/swift/15681.swift b/code_study/Baekjoon/swift/15681.swift new file mode 100644 index 0000000..483e17f --- /dev/null +++ b/code_study/Baekjoon/swift/15681.swift @@ -0,0 +1,66 @@ +struct Tree { + var nodeCount: Int + var root: Int + var subTreeSizes: [Int] + + private var graph: [[Int]] + private var visited: [Bool] + + + init(root R: Int, nodeCount N: Int) { + self.nodeCount = N + self.root = R + + self.graph = Array(repeating: [], count: N+1) + self.visited = Array(repeating: false, count: N+1) + self.subTreeSizes = Array(repeating: 0, count: N+1) + } + + mutating func addEdge(_ u: Int, _ v: Int) { + self.graph[u].append(v) + self.graph[v].append(u) + } + + mutating func calculateSubTreeSize() { + self.dfs(currentNode: self.root) + } + + private mutating func dfs(currentNode: Int) { + self.visited[currentNode] = true + self.subTreeSizes[currentNode] = 1 + + for nextNode in self.graph[currentNode] { + if !self.visited[nextNode] { + self.dfs(currentNode: nextNode) + self.subTreeSizes[currentNode] += self.subTreeSizes[nextNode] + } + } + } +} + +if let input = readLine(), + let info = input.split(separator: " ").compactMap({Int($0)}) as? [Int], + info.count == 3 +{ + let (N, R, Q) = (info[0], info[1], info[2]) + + var tree = Tree(root: R, nodeCount: N) + + for _ in 0..