func processQueries(c int, connections [][]int, queries [][]int) []int {
isOnline := make([]bool, c + 1)
for i := range isOnline {
isOnline[i] = true
}
// build graph
g := make([][]int, c + 1)
for _, v := range connections {
g[v[0]] = append(g[v[0]], v[1])
g[v[1]] = append(g[v[1]], v[0])
}
// identify grids
station2grid := make(map[int]int)
grids := []*grid{}
visited := make([]bool, c + 1)
for stationID := 1; stationID <= c; stationID++ {
if !visited[stationID] {
stations := []int{}
gridID := len(grids)
dfs(stationID, gridID, &stations, visited, g, station2grid)
slices.Sort(stations)
newGrid := grid{stations, 0}
grids = append(grids, &newGrid)
}
}
// run queries
res := []int{}
for _, v := range queries {
if v[0] == 1 {
stationID := v[1]
if isOnline[stationID] {
res = append(res, stationID)
} else { // target station is offline, need to find smallest that is online
gridID := station2grid[stationID]
grid := grids[gridID]
for grid.smallestID < len(grid.stations) && !isOnline[grid.stations[grid.smallestID]] {
grid.smallestID++
}
if grid.smallestID < len(grid.stations) {
res = append(res, grid.stations[grid.smallestID])
} else {
res = append(res, -1)
}
}
} else if v[0] == 2 {
isOnline[v[1]] = false
}
}
return res
}
func dfs(stationID, gridID int, stations *[]int, visited []bool, g [][]int, station2grid map[int]int) {
if visited[stationID] {
return
}
visited[stationID] = true
station2grid[stationID] = gridID
*stations = append(*stations, stationID)
for _, nei := range g[stationID] {
dfs(nei, gridID, stations, visited, g, station2grid)
}
}
type grid struct {
stations []int // stations ordered by id
smallestID int
}