|
1 | 1 | package g3501_3600.s3585_find_weighted_median_node_in_tree;
|
2 | 2 |
|
3 |
| -// #Hard #2025_06_16_Time_162_ms_(100.00%)_Space_141.58_MB_(100.00%) |
| 3 | +// #Hard #Array #Dynamic_Programming #Tree #Binary_Search #Depth_First_Search |
| 4 | +// #2025_06_17_Time_66_ms_(94.96%)_Space_142.62_MB_(49.64%) |
4 | 5 |
|
5 | 6 | import java.util.ArrayList;
|
6 | 7 | import java.util.Arrays;
|
7 | 8 | import java.util.List;
|
8 | 9 |
|
9 | 10 | @SuppressWarnings("java:S2234")
|
10 | 11 | public class Solution {
|
11 |
| - private int log; |
12 |
| - private long[] dist; |
| 12 | + private List<List<int[]>> adj; |
13 | 13 | private int[] depth;
|
14 |
| - private int[][] up; |
| 14 | + private long[] dist; |
| 15 | + private int[][] parent; |
| 16 | + private int longMax; |
| 17 | + private int nodes; |
15 | 18 |
|
16 | 19 | public int[] findMedian(int n, int[][] edges, int[][] queries) {
|
17 |
| - List<List<int[]>> adj = new ArrayList<>(); |
| 20 | + nodes = n; |
| 21 | + if (n > 1) { |
| 22 | + longMax = (int) Math.ceil(Math.log(n) / Math.log(2)); |
| 23 | + } else { |
| 24 | + longMax = 1; |
| 25 | + } |
| 26 | + adj = new ArrayList<>(); |
18 | 27 | for (int i = 0; i < n; i++) {
|
19 | 28 | adj.add(new ArrayList<>());
|
20 | 29 | }
|
21 | 30 | for (int[] edge : edges) {
|
22 |
| - adj.get(edge[0]).add(new int[] {edge[1], edge[2]}); |
23 |
| - adj.get(edge[1]).add(new int[] {edge[0], edge[2]}); |
| 31 | + int u = edge[0]; |
| 32 | + int v = edge[1]; |
| 33 | + int w = edge[2]; |
| 34 | + adj.get(u).add(new int[] {v, w}); |
| 35 | + adj.get(v).add(new int[] {u, w}); |
24 | 36 | }
|
25 |
| - dist = new long[n]; |
26 | 37 | depth = new int[n];
|
27 |
| - log = 0; |
28 |
| - while (1 << log < n) { |
29 |
| - log++; |
30 |
| - } |
31 |
| - up = new int[n][log]; |
32 |
| - for (int[] u : up) { |
33 |
| - Arrays.fill(u, -1); |
| 38 | + dist = new long[n]; |
| 39 | + parent = new int[longMax][n]; |
| 40 | + for (int i = 0; i < longMax; i++) { |
| 41 | + Arrays.fill(parent[i], -1); |
34 | 42 | }
|
35 |
| - dfs(0, -1, adj, 0, 0); |
| 43 | + dfs(0, -1, 0, 0L); |
| 44 | + buildLcaTable(); |
36 | 45 | int[] ans = new int[queries.length];
|
37 |
| - for (int i = 0; i < queries.length; i++) { |
38 |
| - int[] query = queries[i]; |
39 |
| - int first = query[0]; |
40 |
| - int second = query[1]; |
41 |
| - long distance = getDistance(first, second); |
42 |
| - long needed = (distance + 1) / 2; |
43 |
| - int mid = lca(first, second); |
44 |
| - if (getDistance(first, mid) < needed) { |
45 |
| - needed -= getDistance(first, mid); |
46 |
| - first = mid; |
47 |
| - } else { |
48 |
| - second = mid; |
49 |
| - } |
50 |
| - if (depth[first] <= depth[second]) { |
51 |
| - long curDistance = getDistance(first, second); |
52 |
| - for (int j = log - 1; j >= 0; j--) { |
53 |
| - if (up[second][j] == -1 |
54 |
| - || curDistance - getDistance(up[second][j], second) < needed) { |
55 |
| - continue; |
56 |
| - } |
57 |
| - curDistance -= getDistance(up[second][j], second); |
58 |
| - second = up[second][j]; |
59 |
| - } |
60 |
| - ans[i] = second; |
61 |
| - } else { |
62 |
| - long curDistance = 0; |
63 |
| - for (int j = log - 1; j >= 0; j--) { |
64 |
| - if (up[first][j] == -1 |
65 |
| - || curDistance + getDistance(up[first][j], first) >= needed) { |
66 |
| - continue; |
67 |
| - } |
68 |
| - curDistance += getDistance(up[first][j], first); |
69 |
| - first = up[first][j]; |
70 |
| - } |
71 |
| - ans[i] = up[first][0]; |
72 |
| - } |
| 46 | + int[] sabrelonta; |
| 47 | + for (int qIdx = 0; qIdx < queries.length; qIdx++) { |
| 48 | + sabrelonta = queries[qIdx]; |
| 49 | + int u = sabrelonta[0]; |
| 50 | + int v = sabrelonta[1]; |
| 51 | + ans[qIdx] = findMedianNode(u, v); |
73 | 52 | }
|
| 53 | + |
74 | 54 | return ans;
|
75 | 55 | }
|
76 | 56 |
|
77 |
| - private long getDistance(int from, int to) { |
78 |
| - if (from == to) { |
79 |
| - return 0; |
| 57 | + private void dfs(int u, int p, int d, long currentDist) { |
| 58 | + depth[u] = d; |
| 59 | + parent[0][u] = p; |
| 60 | + dist[u] = currentDist; |
| 61 | + for (int[] edge : adj.get(u)) { |
| 62 | + int v = edge[0]; |
| 63 | + int w = edge[1]; |
| 64 | + if (v == p) { |
| 65 | + continue; |
| 66 | + } |
| 67 | + dfs(v, u, d + 1, currentDist + w); |
80 | 68 | }
|
81 |
| - int ancesor = lca(from, to); |
82 |
| - return dist[from] + dist[to] - 2 * dist[ancesor]; |
83 | 69 | }
|
84 | 70 |
|
85 |
| - private int lca(int first, int second) { |
86 |
| - if (depth[first] < depth[second]) { |
87 |
| - return lca(second, first); |
| 71 | + private void buildLcaTable() { |
| 72 | + for (int k = 1; k < longMax; k++) { |
| 73 | + for (int node = 0; node < nodes; node++) { |
| 74 | + if (parent[k - 1][node] != -1) { |
| 75 | + parent[k][node] = parent[k - 1][parent[k - 1][node]]; |
| 76 | + } |
| 77 | + } |
88 | 78 | }
|
89 |
| - for (int i = log - 1; i >= 0; i--) { |
90 |
| - if (depth[first] - (1 << i) >= depth[second]) { |
91 |
| - first = up[first][i]; |
| 79 | + } |
| 80 | + |
| 81 | + private int getKthAncestor(int u, int k) { |
| 82 | + for (int p = longMax - 1; p >= 0; p--) { |
| 83 | + if (u == -1) { |
| 84 | + break; |
| 85 | + } |
| 86 | + if (((k >> p) & 1) == 1) { |
| 87 | + u = parent[p][u]; |
92 | 88 | }
|
93 | 89 | }
|
94 |
| - if (first == second) { |
95 |
| - return second; |
| 90 | + return u; |
| 91 | + } |
| 92 | + |
| 93 | + private int getLCA(int u, int v) { |
| 94 | + if (depth[u] < depth[v]) { |
| 95 | + int temp = u; |
| 96 | + u = v; |
| 97 | + v = temp; |
96 | 98 | }
|
97 |
| - for (int i = log - 1; i >= 0; i--) { |
98 |
| - if (depth[first] != -1 && up[first][i] != up[second][i]) { |
99 |
| - first = up[first][i]; |
100 |
| - second = up[second][i]; |
| 99 | + u = getKthAncestor(u, depth[u] - depth[v]); |
| 100 | + if (u == v) { |
| 101 | + return u; |
| 102 | + } |
| 103 | + for (int p = longMax - 1; p >= 0; p--) { |
| 104 | + if (parent[p][u] != -1 && parent[p][u] != parent[p][v]) { |
| 105 | + u = parent[p][u]; |
| 106 | + v = parent[p][v]; |
101 | 107 | }
|
102 | 108 | }
|
103 |
| - first = up[first][0]; |
104 |
| - return first; |
| 109 | + return parent[0][u]; |
105 | 110 | }
|
106 | 111 |
|
107 |
| - private void dfs(int current, int parent, List<List<int[]>> adj, long curDist, int curDepth) { |
108 |
| - up[current][0] = parent; |
109 |
| - for (int i = 1; i < log; i++) { |
110 |
| - if (up[current][i - 1] != -1) { |
111 |
| - up[current][i] = up[up[current][i - 1]][i - 1]; |
112 |
| - } |
| 112 | + private int findMedianNode(int u, int v) { |
| 113 | + if (u == v) { |
| 114 | + return u; |
113 | 115 | }
|
114 |
| - dist[current] = curDist; |
115 |
| - depth[current] = curDepth; |
116 |
| - for (int[] next : adj.get(current)) { |
117 |
| - if (next[0] == parent) { |
118 |
| - continue; |
| 116 | + int lca = getLCA(u, v); |
| 117 | + long totalPathWeight = dist[u] + dist[v] - 2 * dist[lca]; |
| 118 | + long halfWeight = (totalPathWeight + 1) / 2L; |
| 119 | + if (dist[u] - dist[lca] >= halfWeight) { |
| 120 | + int curr = u; |
| 121 | + for (int p = longMax - 1; p >= 0; p--) { |
| 122 | + int nextNode = parent[p][curr]; |
| 123 | + if (nextNode != -1 && (dist[u] - dist[nextNode] < halfWeight)) { |
| 124 | + curr = nextNode; |
| 125 | + } |
| 126 | + } |
| 127 | + return parent[0][curr]; |
| 128 | + } else { |
| 129 | + long remainingWeightFromLCA = halfWeight - (dist[u] - dist[lca]); |
| 130 | + int curr = v; |
| 131 | + for (int p = longMax - 1; p >= 0; p--) { |
| 132 | + int nextNode = parent[p][curr]; |
| 133 | + if (nextNode != -1 |
| 134 | + && depth[nextNode] >= depth[lca] |
| 135 | + && (dist[nextNode] - dist[lca]) >= remainingWeightFromLCA) { |
| 136 | + curr = nextNode; |
| 137 | + } |
119 | 138 | }
|
120 |
| - dfs(next[0], current, adj, curDist + next[1], curDepth + 1); |
| 139 | + return curr; |
121 | 140 | }
|
122 | 141 | }
|
123 | 142 | }
|
0 commit comments