并查集合并,使用一个数组维护集合个数。
- class Solution {
- static int arr[], cnt[];
-
- public long countPairs(int n, int[][] edges) {
- arr = new int[n];
- cnt = new int[n];
- for (int i = 0; i < n; i++) {
- arr[i] = i;
- cnt[i] = 1;
- }
- for (int i = 0; i < edges.length; i++) {
- int a = edges[i][0];
- int b = edges[i][1];
- a = find(a);
- b = find(b);
- if (a != b) {
- arr[a] = b;
- cnt[b] += cnt[a];
- //cnt[a] = 0;
- }
- }
- long ans = 0, sum = 0;
- for (int i = 0; i < n; i++) {
- if (arr[i] == i) {
- ans += cnt[i] * sum;
- sum += cnt[i];
- }
- }
- return ans;
- }
-
- private int find(int a) {
- if (a != arr[a]) {
- arr[a] = find(arr[a]);
- }
- return arr[a];
- }
- }