LoginSignup
0
1

More than 3 years have passed since last update.

union-find木をGoで実装して理解する

Last updated at Posted at 2020-10-08

初投稿です。

間違っているところがあったら指摘ください、できるだけ直していきたいと思います。

union-find木とは何か?

  • 集合を表現する手法の一つ。 グループ分けを木構造で表現するデータ構造。
    同じ集合に属するということは同じ木に属する、というのがポイント。
    wikipedia

なぜunion-findと呼ばれているか

  • Find: 特定の要素がどの集合に属しているかを求める。
  • Union: 2つの集合を1つにマージする

これらの2つの操作をサポートしているためにUnion-Findと呼ばれる。

wikipediaより

今回の例題

AtCoderBeginnerContest177-D
要約すると
AとBはお友達 という情報がいくつかある。A-B, B-C がお友達だと、A-Cもお友達ということになる。

以上を踏まえて、「一人もお友達がいない」ようにN人の人をグループ分けしたい。

なんグループに分ける必要があるかな?という問題だ。

戦法

僕が最初思いついた解法はこちら。

A, B, C, D, E, F, G, H, Iさんの合計9人がいたとして、

A B
B C
D B
E F
G H
H I

というふうな友人関係だとすると

A-B-C-D
E-F
G-H-I

という3つのお友達グループができる。

これを一人も友達がいないグループに分ける、ということは、直感的に

A B C D
E F
G H I

このような表にして、縦の列でグループを作ると友達が一人もいないグループが出来上がるのがわかるのではないだろうか。

つまり

一番多いお友達グループの人数の分だけ分割すれば、一人も友達がいないグループが出来上がる

ということになる。

しかし、僕の場合ここで問題が発生した。

どうやって集合を表現したらいいのか?

自分はぱっと思いつかなかった。

最初に思いついた手法

集合を表現するスライスを持ったスライスのようなものを思いついた。

だが、これには問題があった。

すべてのリストにAという人が含まれていないかをチェックして、含まれていればそのグループにBという人を追加。
含まれていなければ新たにリストを追加してそこにグループを表現するリストを追加

これをすると、追加の度に最大でリストの要素分の計算量O(N)が発生してしまう。
それを最大で2*10^5人分やらねばならないので、最大O(1+2+3...2*10^5)という計算量になり、O(20,000,100,000=2*10^10)くらいになってしまうため、現在の計算機の能力では2秒以内に終わらせるのは不可能だ。

ではどうすればいいか?

UnionFindの出番である。

UnionFind木の原理

1, 2, 3, 4, 5, 6, 7, 8, 9さんの合計9人がいたとして、

1 2
2 3
4 2
5 6
7 8
8 9

つまり友人関係は

1-2-3-4
5-6
7-8-9

パーツの考察

ベースの木構造の表現

要素に親を持ったスライスを考える。

N=9
parent := make([]int, N)
parent[2]=1

これは2の親が1であるということを表している。

parent[1] == 1

これは1が根だということを表している。

NewUnionFind(num int)

UnionFindのコンストラクタ的なもの。

最初は自分が根であるように初期化しておく。

type UnionFind struct {
    Parent []int
}

func NewUnionFind(num int) UnionFind {
    parent := make([]int, num+1)
    for i := 0; i < len(parent); i++ {
        parent[i] = i
    }
    return UnionFind{parent}
}

Root(x int)

xの根を返す関数

ついでにすべての要素を探索する時に親を更新している

func (uf *UnionFind) Root(x int) int {
    if uf.Parent[x] == x {
        return x
    }
    uf.Parent[x] = uf.Root(uf.Parent[x])
    return uf.Parent[x]
}

Unite(x, y int)

2つの木をマージする関数

同じ木の要素かを調べるのにも使える

func (uf *UnionFind) Unite(x, y int) bool {
    xRoot := uf.Root(x)
    yRoot := uf.Root(y)
    if xRoot == yRoot {
        return false
    }
    uf.Parent[xRoot] = yRoot
    return true
}

上の関数群を使ったメイン関数

func main() {
    scn := NewScanner()
    n, m := scn.Geti(), scn.Geti()
    if m == 0 {
        fmt.Println(1)
        return
    }
    unionFind := NewUnionFind(n)
    for i := 0; i < m; i++ {
        a, b := scn.Geti(), scn.Geti()
                // aとbの属する集合をマージしている
        unionFind.Unite(a, b)
    }

    ar := make([]int, n+1)
    for _, in := range unionFind.Parent {
        ar[unionFind.Root(in)]++
    }
    ans := 0
    for _, in := range ar {
        ans = Max(ans, in)
    }
    fmt.Println(ans)
}

上のコードがどう動くか

fmt.Println(unionFind)を途中に挟んでデバッグしてみるとこういう状態が見える。

9 6
1 2
2 3
4 2
5 6
7 8
8 9
{[0 2 3 3 3 6 6 8 9 9]}
4

すべての処理が終わった後、unionFind.Parentは

[0 2 3 3 3 6 6 8 9 9]

になっているということである

これはどういうことか、少し冗長になるが説明すると

index 0 1 2 3 4 5 6 7 8 9
parent 0 2 3 3 3 6 6 8 9 9

index=1 つまりNo1の人の親はNo2の人

index=2 つまりNo2の人の親はNo3の人

以下同様に全て親を表している

ちなみに、数え上げるときの計算量を抑えるためにRoot()では親の値を根の値で更新している。

しかし、それでも取りこぼしが発生するので、数え上げのときは以下のようにRoot関数を使い親をもう一度探索している。

ar := make([]int, n+1)
for _, in := range unionFind.Parent {
    ar[unionFind.Root(in)]++
}

まとめ

以上のアルゴリズムにより、この問題をunionFind方式で回答することができた。

わかりにくいところや間違っている所の指摘や、こうしたほうがもっとスマートだ、などの意見があればコメントいただけると幸いだ。

最後にフルのソースを載せたいと思う。

package main

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

func main() {
    scn := NewScanner()
    n, m := scn.Geti(), scn.Geti()
    if m == 0 {
        fmt.Println(1)
        return
    }
    unionFind := NewUnionFind(n)
    for i := 0; i < m; i++ {
        a, b := scn.Geti(), scn.Geti()
        unionFind.Unite(a, b)
    }

    ar := make([]int, n+1)
    for _, in := range unionFind.Parent {
        ar[unionFind.Root(in)]++
    }
    ans := 0
    for _, in := range ar {
        ans = Max(ans, in)
    }
    fmt.Println(ans)
}

func Max(a, b int) int {
    if a >= b {
        return a
    }
    return b
}

type UnionFind struct {
    Parent []int
}

func NewUnionFind(num int) UnionFind {
    parent := make([]int, num+1)
    for i := 0; i < len(parent); i++ {
        parent[i] = i
    }
    return UnionFind{parent}
}

func (uf *UnionFind) Root(x int) int {
    if uf.Parent[x] == x {
        return x
    }
    uf.Parent[x] = uf.Root(uf.Parent[x])
    return uf.Parent[x]
}

func (uf *UnionFind) Unite(x, y int) bool {
    xRoot := uf.Root(x)
    yRoot := uf.Root(y)
    if xRoot == yRoot {
        return false
    }
    uf.Parent[xRoot] = yRoot
    return true
}

// Scanner is a base structure
type Scanner struct {
    bufScanner *bufio.Scanner
}

// New inits a scanner
func NewScanner() *Scanner {
    scanner := bufio.NewScanner(os.Stdin)
    scanner.Split(bufio.ScanWords)
    scanner.Buffer(nil, 100000000)
    return &Scanner{scanner}
}

// Geti scans number from stdin
func (scanner *Scanner) Geti() int {
    sc := scanner.bufScanner
    sc.Scan()
    str := sc.Text()

    num, err := strconv.Atoi(str)
    if err != nil {
        panic(err)
    }
    return num
}

0
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
1