20251105 baekjoon

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..<N-1 {
+        if let line = readLine(),
+           let edge = line.split(separator: " ").compactMap({Int($0)}) as? [Int],
+           edge.count == 2
+        {
+            let (u, v) = (edge[0], edge[1])
+            tree.addEdge(u, v)
+        }
+    }
+    
+    tree.calculateSubTreeSize()
+    
+    for _ in 0..<Q {
+        if let q = Int(readLine() ?? "") {
+            print(tree.subTreeSizes[q])
+        }
+    }
+}