AtCoder ABC 126 D - Even Relation / E - 1 or 2 失敗例と解き直し (Go)

irisruneです。今回のコンテストは3完で冷えましたが…Unratedでしたね。

どちらもコンテスト中には解けなかったので解説ACです（正確にはE問題の解説を読んでD問題も解きました）。

ABC 126 D - Even Relation

atcoder.jp

のダイクストラ法で解けるような最短経路問題ばかりやってるとこういう時に躓きますね。まずは失敗例から。

package main

import (
    "fmt"
)

const INF = 10000000000

func main() {
    var n int
    fmt.Scan(&n)

    adj := make([][]int, n)
    dig := make([]int, n)
    visited := make([]bool, n)
    color := make([]int, n)
    dist := make([]int, n)
    queue := make([]int, n)
    for i := range adj {
        line := make([]int, n)
        for j := range line {
            line[j] = INF
        }
        adj[i] = line
        dig[i] = 0
        visited[i] = false
        color[i] = 0
        dist[i] = 0
        queue[i] = 0
    }

    for i := 0; i < n-1; i++ {
        var u, v, w int
        fmt.Scan(&u, &v, &w)
        adj[u-1][v-1] = w
        adj[v-1][u-1] = w
        dig[u-1]++
        dig[v-1]++
    }

    iQueue := -1
    for i := 0; i < n; i++ {
        if dig[i] > 1 {
            continue
        }
        visited[i] = true
        queue[0] = i
        iQueue++
        break
    }

    for iQueue < n-1 {
        iQueueNext := iQueue + 1
        u := queue[iQueue]
        for v, w := range adj[u] {
            if w < INF && !visited[v] {
                dist[v] = dist[u] + w
                if dist[v]%2 == 1 {
                    color[v] = 1
                }
                visited[v] = true
                iQueue++
                queue[iQueue] = v
            }
        }
        iQueue = iQueueNext
    }

    for _, c := range color {
        fmt.Println(c)
    }
}

解き方の大筋としては解説にある通り、適当な頂点からの距離が偶数の点と奇数の点で色分けすればいいです。ただ問題は距離を求めるアルゴリズムの実装ですね。

まず、木構造について葉になる点を基準点にしなければいけないという思い込みがあったために次数の算出と次数1の点の探査という無駄な処理を実装しています。 そして隣接行列を用いてしまっているので計算量メモリ共に無駄が多く、の場合はMLEが出てしまいます。

さらにダイクストラ法を上手く適用できなさそうということに焦ってしまい、スタックではなくキューを使ってDFSを実装しようとしています。当然ながらWAは出ましたし、なんならREまで出てよくわからないことになっていました。

そんな散々な回でしたが、DFSをしっかり実装すれば問題なく解くことができました。また、隣接行列ではなく隣接リストを実装することでMLEも解決しています。

package main

import (
    "fmt"
)

type edge struct {
    tail   int
    weight int
}

func dfs(adj [][]edge, visited []bool, dist []int, color []int, v int) {
    for _, e := range adj[v] {
        if visited[e.tail] {
            continue
        }
        visited[e.tail] = true
        dist[e.tail] = dist[v] + e.weight
        if dist[e.tail]%2 == 1 {
            color[e.tail] = 1
        }
        dfs(adj, visited, dist, color, e.tail)
    }
}

func main() {
    var n int
    fmt.Scan(&n)

    adj := make([][]edge, n)
    visited := make([]bool, n)
    dist := make([]int, n)
    color := make([]int, n)
    for i := range adj {
        line := make([]edge, 0)
        adj[i] = line
    }

    for i := 0; i < n-1; i++ {
        var u, v, w int
        fmt.Scan(&u, &v, &w)
        adj[u-1] = append(adj[u-1], edge{v - 1, w})
        adj[v-1] = append(adj[v-1], edge{u - 1, w})
    }

    visited[0] = true
    dfs(adj, visited, dist, color, 0)

    for _, c := range color {
        fmt.Println(c)
    }
}

実行時間が1700ms余りとかなりギリギリなのでやはり入力を改善する必要はありそうです。

コードレビュー

• 隣接リストに枝の終点と重みをセットで入れているが、この辺りもう少しよい書き方がある？
• DFS関数の引数が多くてちょっとわかりにくい気もする。

E - 1 or 2

atcoder.jp

最終的にやること自体はD問題と変わりません、むしろD問題より実装は楽だと思います。まずは失敗例から。

package main

import (
    "fmt"
)

const INF = 10000000000

func main() {
    var n, m int
    fmt.Scan(&n, &m)
    adj := make([][]int, n)
    for i := range adj {
        line := make([]int, 0)
        adj[i] = line
    }
    dig := make([]int, n)
    for i := 0; i < m; i++ {
        var x, y, z int
        fmt.Scan(&x, &y, &z)
        adj[x-1] = append(adj[x-1], y-1)
        adj[y-1] = append(adj[y-1], x-1)
        dig[x-1]++
        dig[y-1]++
    }
    visited := make([]bool, n)

    ans := 0
    for i := 0; i < n; i++ {
        if !visited[i] {
            ans++
            visited[i] = true
        }
        for j := 0; j < dig[i]; j++ {
            visited[adj[i][j]] = true
        }
    }
    fmt.Println(ans)
}

とりあえずD問題の失敗例コードと比べると隣接行列を隣接リストに変えたりなど少し洗練されていますね。定数INFとかdig配列とか残骸が残っていますが。

解説にある通り、カード全体をグラフ、入力で与えられたカードの組を枝とみなして連結成分の個数を求めればよいです。そのためにDFSを使えばよいということに思い至らず探索アルゴリズムがうまく実装できていません。どこに問題があるか説明するために反例を挙げると、

4 2
1 4 1
3 4 1

が与えられた場合出力として3を返してしまいます（2が正解）。これはカードを単純に番号順に判定しているためで、3番目のカードについて連結成分に属していないという判定をしてしまい間違った結果になってしまいます。

そんなわけで、DFSで実装したコードがこちらです。

package main

import (
    "fmt"
)

func dfs(adj [][]int, visited []bool, u int) {
    for _, v := range adj[u] {
        if visited[v] {
            continue
        }
        visited[v] = true
        dfs(adj, visited, v)
    }
}

func main() {
    var n, m int
    fmt.Scan(&n, &m)
    adj := make([][]int, n)
    for i := range adj {
        line := make([]int, 0)
        adj[i] = line
    }
    for i := 0; i < m; i++ {
        var x, y, z int
        fmt.Scan(&x, &y, &z)
        adj[x-1] = append(adj[x-1], y-1)
        adj[y-1] = append(adj[y-1], x-1)
    }
    visited := make([]bool, n)

    ans := 0
    for i := 0; i < n; i++ {
        if visited[i] {
            continue
        }
        visited[i] = true
        ans++
        dfs(adj, visited, i)
    }
    fmt.Println(ans)
}

D問題のコピペのようですが実はこちらを先にACしました。

雑記

• 本番中にダイクストラ法のGoライブラリについて調べたりしていましたが使い方もよくわからず活用できませんでしたね。
• DFS実装はできたので進歩はしましたが本番中に思いつく所までいきたいですね。
• 他の社員を交えてAtCoderの問題を解いたり対談したレポートを今週中に記事にする予定です。