[20250204] BOJ / 플래3 / 트리와 쿼리 2 / 권혁준

Author: oncsr

khj20006/202502/04 BOJ P3 트리와 쿼리 2.md

@@ -0,0 +1,154 @@
+```java
+
+import java.util.*;
+import java.io.*;
+
+class Main {
+	
+	// IO field
+	static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
+	static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
+	static StringTokenizer st;
+
+	static void nextLine() throws Exception {st = new StringTokenizer(br.readLine());}
+	static int nextInt() {return Integer.parseInt(st.nextToken());}
+	static long nextLong() {return Long.parseLong(st.nextToken());}
+	static void bwEnd() throws Exception {bw.flush();bw.close();}
+	
+	// Additional field
+	static int N;
+	static int[] dep;
+	static int[][] par;
+	static long[][] dist;
+	static int[][] cnt;
+	static List<int[]>[] V;
+	
+	static void initialize() {
+		
+		V = new ArrayList[N+1];
+		dep = new int[N+1];
+		for(int i=1;i<=N;i++) V[i] = new ArrayList<>();
+		par = new int[N+1][18];
+		dist = new long[N+1][18];
+		cnt = new int[N+1][18];
+		
+	}
+	
+	static void input() throws Exception {
+		
+		N = Integer.parseInt(br.readLine());
+		initialize();
+		
+		for(int i=1;i<N;i++) {
+			nextLine();
+			int a = nextInt(), b = nextInt(), c = nextInt();
+			V[a].add(new int[] {b,c});
+			V[b].add(new int[] {a,c});
+		}
+		
+	}
+	
+	static void dfs(int n, int p, int d) {
+		par[n][0] = p;
+		dep[n] = d;
+		for(int[] x : V[n]) {
+			if(x[0] == p) continue;
+			dist[x[0]][0] = x[1];
+			cnt[x[0]][0] = 1;
+			dfs(x[0], n, d+1);
+		}
+	}
+	
+	static void constructSparseTable() {
+		for(int k=1;k<18;k++) {
+			for(int i=1;i<=N;i++) {
+				par[i][k] = par[par[i][k-1]][k-1];
+				dist[i][k] = dist[par[i][k-1]][k-1] + dist[i][k-1];
+				cnt[i][k] = cnt[par[i][k-1]][k-1] + cnt[i][k-1];
+			}
+		}
+	}
+	
+	static int lca(int u, int v) {
+		while(dep[u] != dep[v]) {
+			if(dep[u] > dep[v]) {
+				int diff = dep[u] - dep[v];
+				for(int i=0;i<18;i++) if((diff & (1<<i)) != 0) u = par[u][i];
+			}
+			else {
+				int diff = dep[v] - dep[u];
+				for(int i=0;i<18;i++) if((diff & (1<<i)) != 0) v = par[v][i];
+			}
+		}
+		while(par[u][0] != par[v][0]) {
+			for(int i=17;i>=0;i--) {
+				if(par[u][i] != par[v][i]) {
+					u = par[u][i];
+					v = par[v][i];
+					break;
+				}
+			}
+		}
+		
+		if(u != v) u = par[u][0];
+		return u;
+		
+	}
+	
+	static long query1(int u, int v) {
+
+		long res = 0;
+		int g = lca(u,v);
+		int diff1 = dep[u] - dep[g], diff2 = dep[v] - dep[g];
+		for(int i=0;i<18;i++) {
+			if((diff1 & (1<<i)) != 0) {res += dist[u][i];u = par[u][i];}
+			if((diff2 & (1<<i)) != 0) {res += dist[v][i];v = par[v][i];}
+		}
+		return res;
+		
+	}
+	
+	static int solve(int u, int v, int k) {
+		k--;
+		for(int i=0;i<18;i++) if((k & (1<<i)) != 0) u = par[u][i];
+		return u;
+	}
+	
+	static int query2(int u, int v, int k) {
+		
+		int g = lca(u,v);
+		int diff = dep[u] - dep[g] + 1;
+		if(k <= diff) return solve(u,g,k);
+		k -= (diff-1);
+		diff = dep[v] - dep[g] + 1;
+		k = diff-k+1;
+		return solve(v,g,k);
+		
+	}
+	
+	static void handleQuery() throws Exception {
+		int Q = Integer.parseInt(br.readLine());
+		while(Q-- > 0) {
+			nextLine();
+			int o = nextInt(), u = nextInt(), v = nextInt();
+			if(o == 1) bw.write(query1(u,v)+"\n");
+			else {
+				int k = nextInt();
+				bw.write(query2(u,v,k)+"\n");
+			}
+		}
+	}
+	
+	public static void main(String[] args) throws Exception {
+
+		input();
+		dfs(1,0,0);
+		constructSparseTable();
+		handleQuery();
+		
+		bwEnd();
+	}
+
+}
+
+```