AtCoder 第二回全国統一プログラミング王決定戦予選 B - Counting of Trees (Go)

irisruneです。例によってコンテストには出られませんでした。

問題

atcoder.jp

アルゴリズムは素直ですが、解法のイメージができるかが問われる問題だと思います。

package main

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

var sc *bufio.Scanner

func nextStr() string {
    sc.Scan()
    return sc.Text()
}

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

const MOD = 998244353

func powInt(x, y int) int {
    ret := 1
    for i := 0; i < y; i++ {
        ret *= x
        ret %= MOD
    }
    return ret
}

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Buffer(make([]byte, 0), 1000000001*3)
    sc.Split(bufio.ScanWords)
    n := nextInt()
    depth := make([]int, n)
    d := nextInt()
    if d != 0 {
        fmt.Println(0)
        return
    }
    depth[d]++
    for i := 1; i < n; i++ {
        d := nextInt()
        depth[d]++
    }
    if depth[0] > 1 {
        fmt.Println(0)
        return
    }
    ans := 1
    dPrev := 1
    for _, d := range depth {
        switch {
        case d == 0:
            dPrev = 0
            continue
        case dPrev == 0:
            fmt.Println(0)
            return
        default:
            ans *= powInt(dPrev, d)
            ans %= MOD
            dPrev = d
        }
    }
    fmt.Println(ans)
}

まず条件を満たす木が存在する場合について考えると、基本的な解法は以下の通りです。

1. 頂点1との距離ごとに頂点数をカウントする。
2. から頂点1との距離の最大値まで順番に見て、以下の数値をそれぞれ掛け合わせる。
   1. （距離の頂点数）の（距離の頂点数）乗した値。
3. 掛け合わせた後の値が求める解になる。

次に条件を満たす木が存在しない場合について考えます。それは以下の場合です。

1. 頂点1と頂点1の距離が0でない。
2. 頂点1と頂点1以外の頂点との距離が0である。
3. 頂点1との距離の頂点が存在し、距離の頂点が存在しない。

このうち3.については、上記の解法のように距離の最大値を記録しておくことで自然と解消されます（提出コードでは個別に条件分岐させています）。残り1.と2.については個別に条件分岐させる必要はありますがそれほど難しくはありません。

以上のようにして解を求めることができます。全体の計算量はとなります。

雑記

• C問題がかなり難しいようなのでD問題が解ければ（解説ACでもある程度わかれば）それを、厳しければ過去問を扱う予定です。
  • 今回の記事を見返してみると解法のとっかかりについてノータッチなので、その辺りの説明をするかもしれません。