[Python3/Java/C++] 一题一解:并查集+有序集合(清晰题解)
方法一:并查集 + 有序集合
我们可以使用并查集(Union-Find)来维护电站之间的连接关系,从而确定每个电站所属的电网。对于每个电网,我们使用有序集合(如 Python 中的 SortedList、Java 中的 TreeSet 或 C++ 中的 std::set)来存储该电网中所有在线的电站编号,以便能够高效地查询和删除电站。
具体步骤如下:
- 初始化并查集,处理所有连接关系,将连接的电站合并到同一个集合中。
- 为每个电网创建一个有序集合,初始时将所有电站编号加入对应电网的集合中。
- 遍历查询列表:
- 对于查询 $[1, x]$,首先找到电站 $x$ 所属的电网根节点,然后检查该电网的有序集合:
- 如果电站 $x$ 在线(存在于集合中),则返回 $x$。
- 否则,返回集合中的最小编号电站(如果集合非空),否则返回 -1。
- 对于查询 $[2, x]$,找到电站 $x$ 所属的电网根节点,并将电站 $x$ 从该电网的有序集合中删除,表示该电站离线。
- 对于查询 $[1, x]$,首先找到电站 $x$ 所属的电网根节点,然后检查该电网的有序集合:
- 最后,返回所有类型为 $[1, x]$ 的查询结果。
###python
class UnionFind:
def __init__(self, n):
self.p = list(range(n))
self.size = [1] * n
def find(self, x):
if self.p[x] != x:
self.p[x] = self.find(self.p[x])
return self.p[x]
def union(self, a, b):
pa, pb = self.find(a), self.find(b)
if pa == pb:
return False
if self.size[pa] > self.size[pb]:
self.p[pb] = pa
self.size[pa] += self.size[pb]
else:
self.p[pa] = pb
self.size[pb] += self.size[pa]
return True
class Solution:
def processQueries(
self, c: int, connections: List[List[int]], queries: List[List[int]]
) -> List[int]:
uf = UnionFind(c + 1)
for u, v in connections:
uf.union(u, v)
st = [SortedList() for _ in range(c + 1)]
for i in range(1, c + 1):
st[uf.find(i)].add(i)
ans = []
for a, x in queries:
root = uf.find(x)
if a == 1:
if x in st[root]:
ans.append(x)
elif len(st[root]):
ans.append(st[root][0])
else:
ans.append(-1)
else:
st[root].discard(x)
return ans
###java
class UnionFind {
private final int[] p;
private final int[] size;
public UnionFind(int n) {
p = new int[n];
size = new int[n];
for (int i = 0; i < n; ++i) {
p[i] = i;
size[i] = 1;
}
}
public int find(int x) {
if (p[x] != x) {
p[x] = find(p[x]);
}
return p[x];
}
public boolean union(int a, int b) {
int pa = find(a), pb = find(b);
if (pa == pb) {
return false;
}
if (size[pa] > size[pb]) {
p[pb] = pa;
size[pa] += size[pb];
} else {
p[pa] = pb;
size[pb] += size[pa];
}
return true;
}
}
class Solution {
public int[] processQueries(int c, int[][] connections, int[][] queries) {
UnionFind uf = new UnionFind(c + 1);
for (int[] e : connections) {
uf.union(e[0], e[1]);
}
TreeSet<Integer>[] st = new TreeSet[c + 1];
Arrays.setAll(st, k -> new TreeSet<>());
for (int i = 1; i <= c; i++) {
int root = uf.find(i);
st[root].add(i);
}
List<Integer> ans = new ArrayList<>();
for (int[] q : queries) {
int a = q[0], x = q[1];
int root = uf.find(x);
if (a == 1) {
if (st[root].contains(x)) {
ans.add(x);
} else if (!st[root].isEmpty()) {
ans.add(st[root].first());
} else {
ans.add(-1);
}
} else {
st[root].remove(x);
}
}
return ans.stream().mapToInt(Integer::intValue).toArray();
}
}
###cpp
class UnionFind {
public:
UnionFind(int n) {
p = vector<int>(n);
size = vector<int>(n, 1);
iota(p.begin(), p.end(), 0);
}
bool unite(int a, int b) {
int pa = find(a), pb = find(b);
if (pa == pb) {
return false;
}
if (size[pa] > size[pb]) {
p[pb] = pa;
size[pa] += size[pb];
} else {
p[pa] = pb;
size[pb] += size[pa];
}
return true;
}
int find(int x) {
if (p[x] != x) {
p[x] = find(p[x]);
}
return p[x];
}
private:
vector<int> p, size;
};
class Solution {
public:
vector<int> processQueries(int c, vector<vector<int>>& connections, vector<vector<int>>& queries) {
UnionFind uf(c + 1);
for (auto& e : connections) {
uf.unite(e[0], e[1]);
}
vector<set<int>> st(c + 1);
for (int i = 1; i <= c; i++) {
st[uf.find(i)].insert(i);
}
vector<int> ans;
for (auto& q : queries) {
int a = q[0], x = q[1];
int root = uf.find(x);
if (a == 1) {
if (st[root].count(x)) {
ans.push_back(x);
} else if (!st[root].empty()) {
ans.push_back(*st[root].begin());
} else {
ans.push_back(-1);
}
} else {
st[root].erase(x);
}
}
return ans;
}
};
时间复杂度 $O((c + n + q) \log c)$,空间复杂度 $O(c)$。其中 $c$ 是电站数量,而 $n$ 和 $q$ 分别是连接数量和查询数量。
有任何问题,欢迎评论区交流,欢迎评论区提供其它解题思路(代码),也可以点个赞支持一下作者哈😄~