union-find.hs

-- | Our Union Find Map, this map takes the place of a sparse matrix in
-- implementing a very simple quick union (without weighting or path
-- compression).
type UFM = Map Tag Tag

-- | Find the partition root of a Tag.
findRoot :: Tag -> UFM -> Tag
findRoot tag m =
    -- Chase pointers to the root of a partition, the end is signified by a
    -- link to itself.
    case Map.lookup tag m of
        Just next -> if next == tag then tag else findRoot next m
        Nothing   -> tag

-- | Connect two tag's respective partitions.
union :: Tag -> Tag -> UFM -> UFM
union tag0 tag1 m =
    -- Find the respective roots of the tags.
    let (root0, root1) = (findRoot tag0 m, findRoot tag1 m)
    -- If the roots are not equal, we insert a link into the map.
    --
    -- If this is a self-link, i.e. tag0 == tag1, then we insert a link (from
    -- tag to itself) to ensure that these singleton groups don't get lost.
    --
    -- If the roots are equal, they are in the same partition and there is
    -- nothing to do.
    in if tag0 == tag1 || root0 /= root1
        then Map.insert root0 root1 m else m