利用interface{}在Go中实现通用的DisjointSets

package main

import "fmt"

// DisjointSets 结构体定义
// ranks 存储每个元素的秩（用于优化Union操作）
// p 存储每个元素的父节点（代表元素）
type DisjointSets struct {
    ranks map[interface{}]int64
    p     map[interface{}]interface{}
}

// NewDisjointSets 创建并返回一个新的DisjointSets实例
func NewDisjointSets() *DisjointSets {
    d := DisjointSets{
        ranks: make(map[interface{}]int64),
        p:     make(map[interface{}]interface{}),
    }
    return &d
}

// MakeSet 将元素x添加到不相交集中，作为其自身集合的代表
func (d *DisjointSets) MakeSet(x interface{}) {
    // 检查元素是否已存在，避免重复添加
    if _, exists := d.p[x]; !exists {
        d.p[x] = x
        d.ranks[x] = 0
    }
}

// Link 根据秩（rank）将两个集合的代表元素x和y连接起来
// 秩较小的集合的根节点指向秩较大的集合的根节点
// 如果秩相同，则任意选择一个作为根，并将其秩加1
func (d *DisjointSets) Link(x, y interface{}) {
    if d.ranks[x] > d.ranks[y] {
        d.p[y] = x
    } else {
        d.p[x] = y
        if d.ranks[x] == d.ranks[y] {
            d.ranks[y] += 1
        }
    }
}

// FindSet 查找元素x所属集合的代表元素
// 使用路径压缩优化，将x到根节点路径上的所有节点的父节点直接指向根节点
func (d *DisjointSets) FindSet(x interface{}) interface{} {
    // 如果x不是其自身的父节点，则说明x不是根节点
    if x != d.p[x] {
        // 递归查找根节点，并进行路径压缩
        d.p[x] = d.FindSet(d.p[x])
    }
    return d.p[x]
}

// Union 合并包含元素x和y的两个集合
func (d *DisjointSets) Union(x, y interface{}) {
    // 先找到x和y各自的代表元素，然后将它们连接起来
    d.Link(d.FindSet(x), d.FindSet(y))
}

func main() {
    fmt.Println("--- 整数类型示例 ---")
    dsInt := NewDisjointSets()
    dsInt.MakeSet(1)
    dsInt.MakeSet(2)
    dsInt.MakeSet(3)
    dsInt.MakeSet(4)

    dsInt.Union(1, 2) // 合并 1 和 2
    dsInt.Union(3, 4) // 合并 3 和 4
    dsInt.Union(2, 4) // 合并 2 和 4 (间接合并 1,2,3,4)

    fmt.Printf("元素 1 的代表元素: %v\n", dsInt.FindSet(1)) // 应该与 2,3,4 的代表元素相同
    fmt.Printf("元素 3 的代表元素: %v\n", dsInt.FindSet(3))
    fmt.Printf("元素 2 和 4 是否在同一集合: %t\n", dsInt.FindSet(2) == dsInt.FindSet(4))

    fmt.Println("\n--- 字符串类型示例 ---")
    dsString := NewDisjointSets()
    dsString.MakeSet("apple")
    dsString.MakeSet("banana")
    dsString.MakeSet("cherry")
    dsString.MakeSet("date")

    dsString.Union("apple", "banana")   // 合并 "apple" 和 "banana"
    dsString.Union("cherry", "date")    // 合并 "cherry" 和 "date"
    dsString.Union("banana", "date")    // 合并 "banana" 和 "date" (间接合并 "apple", "banana", "cherry", "date")

    fmt.Printf("元素 'apple' 的代表元素: %v\n", dsString.FindSet("apple"))
    fmt.Printf("元素 'cherry' 的代表元素: %v\n", dsString.FindSet("cherry"))
    fmt.Printf("元素 'apple' 和 'date' 是否在同一集合: %t\n", dsString.FindSet("apple") == dsString.FindSet("date"))

    fmt.Println("\n--- 浮点数类型示例 ---")
    dsFloat := NewDisjointSets()
    dsFloat.MakeSet(1.1)
    dsFloat.MakeSet(2.2)
    dsFloat.MakeSet(3.3)
    dsFloat.Union(1.1, 2.2) // 合并 1.1 和 2.2
    fmt.Printf("元素 1.1 的代表元素: %v\n", dsFloat.FindSet(1.1))
    fmt.Printf("元素 2.2 的代表元素: %v\n", dsFloat.FindSet(2.2))
    fmt.Printf("元素 3.3 的代表元素: %v\n", dsFloat.FindSet(3.3))
}