[LeetCode] 839. Similar String Groups

Problem

Two strings X and Y are similar if we can swap two letters (in different positions) of X, so that it equals Y.

For example, "tars" and "rats" are similar (swapping at positions 0 and 2), and "rats" and "arts" are similar, but "star" is not similar to "tars", "rats", or "arts".

Together, these form two connected groups by similarity: {"tars", "rats", "arts"} and {"star"}. Notice that "tars" and "arts" are in the same group even though they are not similar. Formally, each group is such that a word is in the group if and only if it is similar to at least one other word in the group.

We are given a list A of strings. Every string in A is an anagram of every other string in A. How many groups are there?

Example 1:

Input: ["tars","rats","arts","star"]
Output: 2
Note:

A.length <= 2000
A[i].length <= 1000
A.length * A[i].length <= 20000
All words in A consist of lowercase letters only.
All words in A have the same length and are anagrams of each other.
The judging time limit has been increased for this question.

Solution

class Solution {
    public int numSimilarGroups(String[] A) {
        if (A == null || A.length == 0) return 0;
        int n = A.length;
        UnionFind uf = new UnionFind(n);
        for (int i = 0; i < n-1; i++) {
            for (int j = i+1; j < n; j++) {
                if (canSwap(A[i], A[j])) {
                    uf.union(i, j);
                }
            }
        }
        return uf.size();
    }
    private boolean canSwap(String a, String b) {
        int count = 0;
        for (int i = 0; i < a.length(); i++) {
            if (a.charAt(i) != b.charAt(i) && count++ == 2) return false;
        }
        return true;
    }
}

class UnionFind {
    int[] parents;
    int[] rank;
    int size;
    public UnionFind(int n) {
        parents = new int[n];
        rank = new int[n];
        for (int i = 0; i < n; i++) {
            parents[i] = i;
            rank[i] = 0;
        }
        size = n;
    }
    public int find(int i) {
        if (parents[i] == i) return i;
        int p = find(parents[i]);
        parents[i] = p;
        return p;
    }
    public void union(int a, int b) {
        int pA = find(a);
        int pB = find(b);
        if (pA == pB) return;
        if (rank[pA] > rank[pB]) parents[pB] = pA;
        else if (rank[pB] > rank[pA]) parents[pA] = pB;
        else {
            parents[pA] = pB;
            rank[pB]++;
        }
        size--;
    }
    public int size() {
        return size;
    }
}