Union-Find木について勉強する - どこにでもいるSEの備忘録

まだまだAtCoderをこっそり続けてます。

www.nogawanogawa.work

今回は、Union-Find木について勉強してみたので、そのメモです。

Union-Find木 is なに？
- 例題
試しに解いてみる
- その1
- その2
参考文献
感想

Union-Find木 is なに？

Union-Find木は、データのグループ化を高速に行うことができるデータ構造です。下記のような2種類のクエリ（タスク）を行うことで、実現されます。

Find
- 特定の要素がどの集合に属しているかを判定する（2つの要素が同じ集合に属しているかの判定にも使用される）
Union
- 2つの集合を1つに統合する

詳しくはこちらのスライドの解説が非常にわかりやすいです。

Union find(素集合データ構造) from AtCoder Inc.

例題

上のスライドが解説になった問題について解いてみます。

atcoder.jp

シンプルにUnion-Find木を使ってやればできます。経路圧縮もランクも使用してません。

コード

package main

import (
    "bufio"
    "fmt"
    "os"
    "strconv"
)

type UnionFind struct {
    par []int
}

func newUnionFind(N int) *UnionFind {
    u := new(UnionFind)
    u.par = make([]int, N)
    for i := range u.par {
        u.par[i] = -1
    }
    return u
}

// 根を返す
func (u UnionFind) root(x int) int {
    if u.par[x] < 0 {
        return x // 初期値マイナスのため自分が根
    }
    u.par[x] = u.root(u.par[x]) // 親をたどっていって根を返す
    return u.par[x]
}

// union クエリ
func (u UnionFind) unite(x, y int) {
    xr := u.root(x)
    yr := u.root(y)
    if xr == yr {
        return
    }
    u.par[yr] += u.par[xr]
    u.par[xr] = yr
}

// find クエリ
func (u UnionFind) same(x, y int) bool {
    if u.root(x) == u.root(y) {
        return true
    }
    return false
}

var sc = bufio.NewScanner(os.Stdin)

func nextInt() int {
    sc.Scan()
    i, e := strconv.ParseInt(sc.Text(),10,32)
    
    if e != nil {
        panic(e)
    }

    return int(i)
}

func main() {
    sc.Split(bufio.ScanWords)

    var N, Q int
    fmt.Scanf("%d %d", &N, &Q)

    u := newUnionFind(N)

    P := make([]int, Q)
    A := make([]int, Q)
    B := make([]int, Q)

    for i:=0; i<Q; i++{
        P[i] = nextInt()
        A[i] = nextInt()
        B[i] = nextInt()
    }

    for i:=0; i<Q; i++{
        if P[i] == 0{
            u.unite(A[i], B[i])
        } else {
            if u.same(A[i], B[i]){
                fmt.Printf("Yes\n")
            } else {
                fmt.Printf("No\n")
            }
        }
    }
}

試しに解いてみる

それでは、実際の問題をいくつか解いていきたいと思います。

その1

atcoder.jp

こちらの問題は、一見するとUnion-Find木は関係ないように見えます。しかし、見方を変えればUnion-Find木っぽく見ることができる問題です。

はじめに得られる点をNode、交換可能なペアを辺と考えると、無向グラフ構造が構築できます。その無向グラフに対してindex=iにいるNodeとvalue=iが同じ根を持てばカウントすることができます。

コード

package main

import (
    "bufio"
    "fmt"
    "os"
    "strconv"
)

type UnionFind struct {
    par []int
    val []int
}

func newUnionFind(N int, p []int) *UnionFind {
    u := new(UnionFind)
 
    u.par = make([]int, N+1)
    u.val = make([]int, N+1)
 
    for i:=0; i<len(p); i++ {
        u.val[i] = p[i]
        u.par[i] = -1
    }
 
    return u
}

// 根を返す
func (u UnionFind) root(x int) int {
    if u.par[x] < 0 {
        return x // 初期値マイナスのため自分が根
    }
    u.par[x] = u.root(u.par[x]) // 親をたどっていって根を返す
    return u.par[x]
}

// union クエリ
func (u UnionFind) unite(x, y int) {
    xr := u.root(x)
    yr := u.root(y)
    if xr == yr {
        return
    }
    u.par[yr] += u.par[xr]
    u.par[xr] = yr
}

// find クエリ
func (u UnionFind) same(x, y int) bool {
    if u.root(x) == u.root(y) {
        return true
    }
    return false
}

var sc = bufio.NewScanner(os.Stdin)

func nextInt() int {
    sc.Scan()
    i, e := strconv.ParseInt(sc.Text(),10,32)
    if e != nil {
        panic(e)
    }
    return int(i)
}

func main() {
    sc.Split(bufio.ScanWords)
 
    var N, M int
    fmt.Scanf("%d %d",&N, &M)

    dict := make(map[int]int)

    p := make([]int, N+1)
    for i:=1; i<N+1; i++{
        p[i] = nextInt()
        dict[p[i]] = i
    }
 
    u := newUnionFind(N, p)

    var x, y int

    for i:=0; i<M; i++{
        x = nextInt()    
        y = nextInt()

        u.unite(x,y)
    }

    var count int = 0
    for i:=1; i<N+1; i++{
        if u.root(i) == u.root(dict[i]){
            count ++ 
        }
    }

    fmt.Printf("%d\n", count)
}

その2

atcoder.jp

こちらの問題は、Union-Find木で交差点を結合していき、最終的にすべての交差点の根が何種類あるかをカウントします。カウントした根の種類数-1個だけ道路を作れば、すべての交差点が辺でつながった一つの無向グラフにすることができます。

コード

package main

import (
    "bufio"
    "fmt"
    "os"
    "strconv"
)

type UnionFind struct {
    par []int64
}

func newUnionFind(N int64) *UnionFind {
    u := new(UnionFind)
 
    u.par = make([]int64, N+1)

    var i int64
    for i=0; i<N+1; i++ {
        u.par[i] = -1
    }
 
    return u
}

// 根を返す
func (u UnionFind) root(x int64) int64 {
    if u.par[x] < 0 {
        return x // 初期値マイナスのため自分が根
    }
    u.par[x] = u.root(u.par[x]) // 親をたどっていって根を返す
    return u.par[x]
}

// union クエリ
func (u UnionFind) unite(x, y int64) {
    xr := u.root(x)
    yr := u.root(y)
    if xr == yr {
        return
    }
    u.par[yr] += u.par[xr]
    u.par[xr] = yr
}

// find クエリ
func (u UnionFind) same(x, y int64) bool {
    if u.root(x) == u.root(y) {
        return true
    }
    return false
}

var sc = bufio.NewScanner(os.Stdin)

func nextInt() int64 {
    sc.Scan()
    i, e := strconv.ParseInt(sc.Text(),10,32)
    if e != nil {
        panic(e)
    }
    return i
}

func main() {
    sc.Split(bufio.ScanWords)
 
    var N, M int64
    fmt.Scanf("%d %d",&N, &M)

    var i int64
 
    u := newUnionFind(N)

    var a, b int64

    for i=0; i<M; i++{
        a = nextInt()    
        b = nextInt()

        u.unite(a,b)
    }

    dict := make(map[int64]int64)
    var count int = 0
    for i=1; i<N+1; i++{

        _, ok := dict[u.root(i)]
        if ok==false{
            dict[u.root(i)] = i
            count ++
        }
    }

    fmt.Printf("%d\n", count-1)
}