問題概要

SNSに $N$ 人の人がいます。この $N$ 人の間には、 $M$ 組の「友達関係」と、 $K$ 組の「ブロック関係」が存在します(いずれも双方向)。各人において、友達の友達の…とたどって行き着く人のうち、直接の友達でもブロック関係にもない人(友達候補)の数を求めてください。

atcoder.jp

前置き

この問題はUnionFindを使う解法が最も多そうです。 UnionFindについてはこれがとても分かりやすかったです。

qiita.com

ただ自分はUnionFindを知らなかったため、幅優先探索(BFS)で解きました。

(UnionFindでの実装は最後に載せてあります。)

各連結成分を求めてあげて、その中で自分と友達でもなく、ブロック関係でもない人の数を求めれば良いので、

(人 $i$ の友達候補の数) = (人 $i$ の連結成分のサイズ) − (人 $i$ と人 $j$ が同じ連結成分に含まれて、人 $i$ と人 $j$ が友達関係もしくはブロック関係にあるような $j$ の数) −1 (自分自身)

となります。

解法概要

まず、友達関係の入力から無向グラフを作成します。

その後、ノード $i$ の連結成分を求めるため、ノード $i$ を始点としたBFSによって隣接しているノードを探索していきます。もしも自分が含まれている連結成分が、同じ連結成分の他のノードを始点としたBFSによってすでに求まっている場合は、それを再利用します。

データ構造としては、 sets map[*Node]map[*Node]bool という、ハッシュを入れ子にした形を使います。 sets[people[i]] とすると、人 $i$ の連結成分の要素をハッシュとして取り出せます。

また、人 $i$ の直接の友達は、ノード $i$ の隣接ノードなので、その数はすぐ求まります。

さらに、ブロック関係については、人 $i$ と同じ連結成分に含まれているとき(sets[people[i]]の要素であるとき)のみ、数をカウントしていきます。

以上によって必要な要素は求まります。

コード

言語はGoです。Goにはキューが標準ライブラリに用意されていないため、自分で実装しています(とは言っても簡単にできます)。

findSetsという関数によって、グラフの連結成分の集合(sets)を求めている部分が今回の肝です。

package main

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

type Queue struct {
    data []interface{}
}

func (q *Queue) Push(v interface{}) {
    q.data = append(q.data, v)
}

func (q *Queue) Pop() (interface{}, error) {
    if q.Empty() {
        return nil, fmt.Errorf("queue is empty")
    }
    v := q.data[0]
    q.data = q.data[1:]
    return v, nil
}

func (q *Queue) Empty() bool {
    return len(q.data) == 0
}

func NewQueue() *Queue {
    return &Queue{make([]interface{}, 0)}
}

type Graph struct {
    Nodes []*Node
}

type Node struct {
    Val    interface{}
    Neighs []*Node
}

func dfs(g *Graph, node *Node, groups map[*Node]*Node) map[*Node]bool {
    set := make(map[*Node]bool)
    q := NewQueue()
    q.Push(node)
    set[node] = true
    for !q.Empty() {
        src, _ := q.Pop()
        // Queueの要素がinterface{}なので型キャストが必要
        if v, ok := src.(*Node); ok {
            for _, neigh := range v.Neighs {
                if !set[neigh] {
                    q.Push(neigh)
                    set[neigh] = true
                    groups[neigh] = node
                }
            }
        }
    }
    return set
}

func findSets(g *Graph) map[*Node]map[*Node]bool {
    sets := make(map[*Node]map[*Node]bool)
    groups := make(map[*Node]*Node)
    nodes := g.Nodes
    for _, node := range nodes {
        // 同じ連結成分の他のノードによって求まっている場合、それを再利用
        if leader, ok := groups[node]; ok {
            sets[node] = sets[leader]
            continue
        }
        sets[node] = dfs(g, node, groups)
    }
    return sets
}

var sc = bufio.NewScanner(os.Stdin)

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

func main() {
    sc.Split(bufio.ScanWords)
    n, m, k := nextInt(), nextInt(), nextInt()

    people := make([]*Node, n)
    for i := 0; i < n; i++ {
        people[i] = &Node{i, make([]*Node, 0)}
    }

    for i := 0; i < m; i++ {
        a, b := nextInt()-1, nextInt()-1

        people[a].Neighs = append(people[a].Neighs, people[b])
        people[b].Neighs = append(people[b].Neighs, people[a])
    }

    g := &Graph{people}

    // 連結成分を求める
    sets := findSets(g)

    blocked := make(map[*Node]int)
    for i := 0; i < k; i++ {
        c, d := nextInt()-1, nextInt()-1

        setc := sets[people[c]]
        if setc[people[d]] {
            // 同じ連結成分に含まれるときだけカウント
            blocked[people[c]]++
        }

        setd := sets[people[d]]
        if setd[people[c]] {
            // 同じ連結成分に含まれるときだけカウント
            blocked[people[d]]++
        }
    }

    for i := 0; i < n; i++ {
        pi := people[i]
        set := sets[pi]
        // 上の式そのまま
        fmt.Printf("%v ", len(set)-len(pi.Neighs)-blocked[pi]-1)
    }
    fmt.Println()
}

感想

BFSによって連結成分を求める、という部分は他の問題にも応用できそう？
最初、入力をfmt.Scan()で受けていたら、そこがボトルネックになってTLEになっていたので注意…ちゃんとbufioを使おう
グラフとBFSの良い練習になりました

おまけ

勉強がてらUnionFindでも実装してみたので、載せておきます。

package main

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

type UnionFind struct {
    Par   []int
    Ranks []int
    Sizes []int
}

func NewUnionFind(n int) *UnionFind {
    uf := &UnionFind{make([]int, n), make([]int, n), make([]int, n)}
    for i := range uf.Par {
        uf.Par[i] = i
        uf.Ranks[i] = 0
        uf.Sizes[i] = 1
    }
    return uf
}

func (uf *UnionFind) Root(x int) int {
    if x == uf.Par[x] {
        return x
    }
    uf.Par[x] = uf.Root(uf.Par[x]) // 経路圧縮
    return uf.Par[x]
}

func (uf *UnionFind) Unite(x, y int) {
    rx := uf.Root(x)
    ry := uf.Root(y)
    if rx == ry {
        return
    }
    if uf.Ranks[rx] < uf.Ranks[ry] {
        uf.Par[rx] = ry
        uf.Sizes[ry] += uf.Sizes[rx]
    } else {
        uf.Par[ry] = rx
        uf.Sizes[rx] += uf.Sizes[ry]
        if uf.Ranks[rx] == uf.Ranks[ry] {
            uf.Ranks[rx]++
        }
    }
}

func (uf *UnionFind) Same(x, y int) bool {
    return uf.Root(x) == uf.Root(y)
}

func (uf *UnionFind) Size(x int) int {
    return uf.Sizes[uf.Root(x)]
}

var sc = bufio.NewScanner(os.Stdin)

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

func main() {
    sc.Split(bufio.ScanWords)
    n, m, k := nextInt(), nextInt(), nextInt()
    uf := NewUnionFind(n)

    friends := make(map[int]int)
    for i := 0; i < m; i++ {
        a, b := nextInt()-1, nextInt()-1
        uf.Unite(a, b)
        friends[a]++
        friends[b]++
    }

    // Sizesをufに持たせない場合
    // sizes := make([]int, n)
    // for i := range sizes {
    //     r := uf.Root(uf.Par[i])
    //     sizes[r]++
    // }

    blocks := make([]int, n)
    for i := 0; i < k; i++ {
        c, d := nextInt()-1, nextInt()-1
        if uf.Same(c, d) {
            blocks[c]++
            blocks[d]++
        }
    }

    for i := 0; i < n; i++ {
        ans := uf.Size(i) - friends[i] - blocks[i] - 1
        fmt.Printf("%v ", ans)
    }
    fmt.Println()
}

shomanのブログ

ただの備忘録

AtCoder ABC157 D - Friend Suggestions BFSとUnionFind

問題概要

前置き

解法概要

コード

感想

おまけ