|
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