package main

import (
    "fmt"
)

type Label int

type Pos struct {
    x, y int
}

type DisjointSets struct {
    parents map[Pos]Pos
    rank    map[Pos]int
    labels  map[Pos]Label
}

func NewDisjointSets() *DisjointSets {
    return &DisjointSets{
        parents: make(map[Pos]Pos),
        rank:    make(map[Pos]int),
        labels:  make(map[Pos]Label),
    }
}

func (s *DisjointSets) MakeSet(u Pos, label Label) {
    // ignore if set already exists
    if _, ok := s.parents[u]; ok {
        return
    }

    s.labels[u] = label
    s.parents[u] = u
    s.rank[u] = 1
}

func (s *DisjointSets) FindSet(u Pos) Pos {
    assertExists(s.parents, u)

    var path []Pos

    n := u

    for {
        if s.parents[n] == n {
            break
        }

        path = append(path, n)
        n = s.parents[n]
    }

    // path compression
    for _, p := range path {
        s.parents[p] = n
    }

    return n
}

func (s *DisjointSets) Label(u Pos) Label {
    assertExists(s.parents, u)

    parent := s.FindSet(u)
    return s.labels[parent]
}

func (s *DisjointSets) union(u, v Pos) {
    assertExists(s.parents, u, v)

    uParent := s.FindSet(u)
    vParent := s.FindSet(v)

    // already in the same set
    if uParent == vParent {
        return
    }

    uLabel := s.Label(u)
    vLabel := s.Label(v)

    label := minLabel(uLabel, vLabel)

    if s.rank[uParent] == s.rank[vParent] {
        s.rank[uParent]++
        s.parents[vParent] = uParent
        s.labels[uParent] = label
    } else if s.rank[uParent] > s.rank[vParent] {
        s.parents[vParent] = uParent
        s.labels[uParent] = label
    } else {
        s.parents[uParent] = vParent
        s.labels[vParent] = label
    }
}

func (s *DisjointSets) SameComponent(u, v Pos) bool {
    assertExists(s.parents, u, v)

    return s.FindSet(u) == s.FindSet(v)
}

func (s *DisjointSets) Dump() {
    fmt.Println("====================")

    for u, parent := range s.parents {
        fmt.Printf(
            "(%d,%d) parent: (%d,%d) label %dn",
            u.x,
            u.y,
            parent.x,
            parent.y,
            s.Label(u),
        )
    }
}

func minLabel(labels ...Label) Label {
    if len(labels) == 0 {
        panic("len(labels) must be >= 0")
    }

    var min = labels[0]

    for _, label := range labels {
        if label < min {
            min = label
        }
    }

    return min
}

func assertExists(parents map[Pos]Pos, check ...Pos) {
    for _, u := range check {
        if _, ok := parents[u]; !ok {
            panic(fmt.Sprintf("Pos (%d,%d) is not a known set.", u.x, u.y))
        }
    }
}

func main() {
    sets := NewDisjointSets()

    sets.MakeSet(Pos{0, 0}, 1)
    sets.MakeSet(Pos{1, 0}, 2)
    sets.MakeSet(Pos{2, 0}, 2)

    sets.MakeSet(Pos{0, 1}, 1)
    sets.MakeSet(Pos{1, 1}, 3)
    sets.MakeSet(Pos{2, 1}, 2)

    sets.MakeSet(Pos{0, 2}, 3)
    sets.MakeSet(Pos{1, 2}, 3)
    sets.MakeSet(Pos{2, 2}, 4)

    sets.Dump()

    sets.union(Pos{0, 0}, Pos{1, 0})
    sets.union(Pos{0, 0}, Pos{1, 1})

    sets.Dump()

    fmt.Printf("%#vn", sets.SameComponent(Pos{1, 1}, Pos{2, 2}))
    fmt.Printf("%#vn", sets.SameComponent(Pos{1, 1}, Pos{0, 0}))
    fmt.Printf("%#vn", sets.SameComponent(Pos{0, 2}, Pos{0, 2}))

    sets.union(Pos{1, 1}, Pos{1, 1})

    sets.Dump()

    sets.union(Pos{1, 1}, Pos{2, 2})

    sets.Dump()
}