1
+ package com .problems .graph ;
2
+
3
+ import com .util .PrintUtl ;
4
+
5
+ import java .util .ArrayList ;
6
+ import java .util .Arrays ;
7
+ import java .util .List ;
8
+ import java .util .TreeSet ;
9
+
10
+ /*
11
+ * Problem link:
12
+ * https://leetcode.com/problems/power-grid-maintenance/description/
13
+ * Solution link:
14
+ *
15
+ */
16
+ public class PowerGridMaintenance {
17
+ public static void main (String [] args ) {
18
+ type1 ();
19
+ type2 ();
20
+ }
21
+
22
+ private static void type2 () {
23
+ }
24
+
25
+ // Using union find to group the power stations then resolve the queries one by one
26
+ // todo this is not the optimal solution, we can use a better approach
27
+ private static void type1 () {
28
+ int c = 5 ;
29
+ int [][] connections = {{1 , 2 }, {2 , 3 }, {3 , 4 }, {4 , 5 }};
30
+ int [][] queries = {{1 , 3 }, {2 , 1 }, {1 , 1 }, {2 , 2 }, {1 , 2 }};
31
+ int [] ans = processQueries1 (c , connections , queries );
32
+ PrintUtl .print (ans );
33
+ }
34
+
35
+ public static int [] processQueries1 (int c , int [][] connections , int [][] queries ) {
36
+ // creating union find structure
37
+ int [] parent = new int [c + 1 ];
38
+ int [] wt = new int [c + 1 ];
39
+ for (int i = 1 ; i <= c ; i ++) parent [i ] = i ;
40
+ Arrays .fill (wt , 1 );
41
+ buildUnionFind (connections , parent , wt );
42
+
43
+ TreeSet <Integer >[] powerStations = new TreeSet [c + 1 ];
44
+ for (int node = 1 ; node <= c ; node ++) {
45
+ // redoing the union find to get the parent of the node, if the path compression is not done
46
+ int parentNode = parent [node ] = parent (parent , node );
47
+ if (powerStations [parentNode ] == null ) {
48
+ TreeSet <Integer > set = new TreeSet <>();
49
+ set .add (parentNode );
50
+ powerStations [parentNode ] = set ;
51
+ }
52
+ powerStations [parentNode ].add (node );
53
+ }
54
+ // even if the node is removed, we will still need to find the parent node, we will keep a copy of the parent array
55
+ int [] copyParent = Arrays .copyOf (parent , parent .length );
56
+ List <Integer > ans = new ArrayList <>();
57
+ // Process each query
58
+ for (int [] query : queries ) {
59
+ int op = query [0 ], node = query [1 ];
60
+ int parentNode = parent [node ];
61
+ if (op == 1 ) {
62
+ // copyParent[node] != -1 means it is not removed yet, so the ans will be the same node
63
+ if (copyParent [node ] != -1 ) {
64
+ ans .add (node );
65
+ } else {
66
+ // the node is removed, we need to find the parent node
67
+ TreeSet <Integer > set = powerStations [parentNode ];
68
+ int queryAns = !set .isEmpty () ? set .first () : -1 ;
69
+ ans .add (queryAns );
70
+ }
71
+ } else {
72
+ // we will remove the node from the power station and update the copyParent
73
+ powerStations [parentNode ].remove (node );
74
+ copyParent [node ] = -1 ;
75
+ }
76
+ }
77
+ // Convert the list to an array
78
+ return convertListToArray (ans );
79
+ }
80
+
81
+ private static void buildUnionFind (int [][] connections , int [] parent , int [] wt ) {
82
+ for (int [] edge : connections ) {
83
+ int x = edge [0 ], y = edge [1 ];
84
+ int root1 = parent (parent , x );
85
+ int root2 = parent (parent , y );
86
+ if (root1 == root2 ) continue ;
87
+ if (wt [root1 ] >= wt [root2 ]) {
88
+ wt [root1 ] += wt [root2 ];
89
+ wt [root2 ] = 0 ;
90
+ parent [root2 ] = root1 ;
91
+ } else {
92
+ wt [root2 ] += wt [root1 ];
93
+ wt [root1 ] = 0 ;
94
+ parent [root1 ] = root2 ;
95
+ }
96
+ }
97
+ }
98
+
99
+ private static int [] convertListToArray (List <Integer > list ) {
100
+ int [] arr = new int [list .size ()];
101
+ for (int i = 0 ; i < list .size (); i ++) {
102
+ arr [i ] = list .get (i );
103
+ }
104
+ return arr ;
105
+ }
106
+
107
+ static int parent (int [] par , int node ) {
108
+ int pt = node ;
109
+ while (pt != par [pt ]) {
110
+ pt = par [pt ];
111
+ }
112
+ int root = pt ;
113
+ pt = node ;
114
+ while (pt != par [pt ]) {
115
+ int nextPt = par [pt ];
116
+ par [pt ] = root ;
117
+ pt = nextPt ;
118
+ }
119
+ return root ;
120
+ }
121
+ }
0 commit comments