Untitled

import Foundation

// Union-Find Data Structure
class UnionFind {
    private var parent: [Int]
    private var rank: [Int]
    
    init(_ size: Int) {
        parent = Array(0..<size)
        rank = Array(repeating: 1, count: size)
    }
    
    func find(_ p: Int) -> Int {
        if parent[p] != p {
            parent[p] = find(parent[p])
        }
        return parent[p]
    }
    
    func union(_ p: Int, _ q: Int) {
        let rootP = find(p)
        let rootQ = find(q)
        
        if rootP != rootQ {
            if rank[rootP] > rank[rootQ] {
                parent[rootQ] = rootP
            } else if rank[rootP] < rank[rootQ] {
                parent[rootP] = rootQ
            } else {
                parent[rootQ] = rootP
                rank[rootP] += 1
            }
        }
    }
}

// Helper function to check if two strings are similar
func areSimilar(_ s1: String, _ s2: String) -> Bool {
    if s1 == s2 { return true }
    
    var diffIndices = [Int]()
    let chars1 = Array(s1)
    let chars2 = Array(s2)
    
    for i in 0..<chars1.count {
        if chars1[i] != chars2[i] {
            diffIndices.append(i)
        }
        if diffIndices.count > 2 {
            return false
        }
    }
    
    return diffIndices.count == 2 && chars1[diffIndices[0]] == chars2[diffIndices[1]] && chars1[diffIndices[1]] == chars2[diffIndices[0]]
}

// Main function to find the number of groups
func numSimilarGroups(_ strs: [String]) -> Int {
    let n = strs.count
    let uf = UnionFind(n)
    
    for i in 0..<n {
        for j in i + 1..<n {
            if areSimilar(strs[i], strs[j]) {
                uf.union(i, j)
            }
        }
    }
    
    var uniqueGroups = Set<Int>()
    for i in 0..<n {
        uniqueGroups.insert(uf.find(i))
    }
    
    return uniqueGroups.count
}

// Test cases
let testCases = [
    ["tars", "rats", "arts", "star"],
    ["omv", "ovm"],
    ["abc", "bca", "cab", "xyz"]
]

for (index, testCase) in testCases.enumerated() {
    print("Test Case \(index + 1): \(numSimilarGroups(testCase))")
}