あるエンジニアのAtCoder奮闘記

東京都港区にあるアミフィアブル株式会社のエンジニアが、AtCoderで解いた問題について振り返ったりしていく会社公認のブログです。

AtCoder AGC 017 A - Biscuits で色々な解き方を試してみる (Go)

irisruneです。AtCoder ProblemsのRecommendations機能のおかげで精進に困らなくていいですね。

atcoder.jp

失敗例含め4回解いていますが、大筋はあまり変わらない上に想定解法が天才解法なので全部嘘解法です。

失敗例

package main

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

var sc *bufio.Scanner

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

func bico(n, r int) int {
    if n < r {
        return 0
    }
    if r > 0 && n/r < 2 {
        r = n - r
    }
    denom, numer := 1, 1
    for i := 1; i <= r; i++ {
        denom *= n - i + 1
        numer *= i
    }
    return denom / numer
}

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Split(bufio.ScanWords)
    n, p := nextInt(), nextInt()

    cnt0, cnt1 := 0, 0
    for i := 0; i < n; i++ {
        a := nextInt()
        switch {
        case a%2 == 0:
            cnt0++
        default:
            cnt1++
        }
    }

    ans0 := 0
    ans1 := 0
    for j := 0; j <= cnt0; j++ {
        ans0 += bico(cnt0, j)
    }
    for k := p; k <= cnt1; k += 2 {
        ans1 += bico(cnt1, k)
    }

    fmt.Println(ans0 * ans1)
}

前回の雑記で触れた入力の最適化を行っているのでテンプレートが少し長くなっています。

大筋としては、ビスケットの数が偶数の袋と奇数の袋の数をそれぞれ数え上げ、偶数の袋を選ぶパターン数と奇数の袋を選ぶパターン数を二項係数として計算して掛け合わせるというものになります。

剰余を要求されていないため逆元を知らなくても使わなくても解くことができますね。

ただしこのコードは、例えばN=50のとき計算結果(の分母部分)に50\times49\times...\times25\simeq1.96\times10^{29}が入ってしまうので64bit型整数ではオーバーフローしてしまいます。そのためWAとなってしまいました。

ちなみに、計算量はO(N^3)と大きめです。

有理数の形ではなく実数の形で計算結果をもつ

package main

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

var sc *bufio.Scanner

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

func bico(n, r int) int {
    if n < r {
        return 0
    }
    if n < 2*r {
        r = n - r
    }

    ret := 1.0
    for i := 1; i <= r; i++ {
        ret *= float64(n - i + 1)
        ret /= float64(i)
    }
    return int(ret)
}

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Split(bufio.ScanWords)
    n, p := nextInt(), nextInt()

    cnt0, cnt1 := 0, 0
    for i := 0; i < n; i++ {
        a := nextInt()
        switch {
        case a%2 == 0:
            cnt0++
        default:
            cnt1++
        }
    }

    ans0 := 0
    ans1 := 0
    for j := 0; j <= cnt0; j++ {
        ans0 += bico(cnt0, j)
    }
    for k := p; k <= cnt1; k += 2 {
        ans1 += bico(cnt1, k)
    }

    fmt.Println(ans0 * ans1)
}

このコードでは、有理数の形で計算結果を持つとオーバーフローしてしまう問題を、逐一除算を行うことで実数の形にして解決しています。

欠点としては整数型でデータを持てないため、誤差が出る危険性があるかもしれません。二項係数の計算アルゴリズムで誤差が出るかどうかはよくわかりませんが…

また、計算量は先の方法と変わらずO(N^3)となります。

パスカルの三角形を使う

package main

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

var sc *bufio.Scanner

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

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Split(bufio.ScanWords)
    n, p := nextInt(), nextInt()

    cnt0, cnt1 := 0, 0
    for i := 0; i < n; i++ {
        a := nextInt()
        switch {
        case a%2 == 0:
            cnt0++
        default:
            cnt1++
        }
    }

    bico := make([][]int, n+1)
    for i := range bico {
        bicoLine := make([]int, n+1)
        bicoLine[0] = 1
        bicoLine[i] = 1
        bico[i] = bicoLine
    }

    for i := range bico {
        for r := range bico[i] {
            if i >= 1 && r >= 1 && r < i {
                bico[i][r] = bico[i-1][r-1] + bico[i-1][r]
            }
        }
    }

    ans0 := 0
    ans1 := 0
    for j := 0; j <= cnt0; j++ {
        ans0 += bico[cnt0][j]
    }
    for k := p; k <= cnt1; k += 2 {
        ans1 += bico[cnt1][k]
    }

    fmt.Println(ans0 * ans1)
}

パスカルの三角形と呼ばれる関係として、

 \displaystyle
\binom{n}{0} = \binom{n}{n} = 1

 \displaystyle
\binom{n}{k} = \binom{n-1}{k-1} + \binom{n-1}{k}

というものがあります。1これを用いて予めすべての二項係数を計算しておくことで問題を解くことができます。

(知らないと/調べないと辿り着かないこと以外は)欠点らしい欠点はなく、N\leq50の制約条件下ではオーバーフローの心配もなく、計算量もO(N^2)で済むのでよい解き方だと思います。

逐一約分する

package main

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

var sc *bufio.Scanner

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

func gcd(a, b int) int {
    if a < b {
        a, b = b, a
    }
    if a%b == 0 {
        return b
    }
    return gcd(b, a%b)
}

func bico(n, r int) int {
    if n < r {
        return 0
    }
    if r > 0 && n/r < 2 {
        r = n - r
    }
    denom, numer := 1, 1
    for i := 1; i <= r; i++ {
        denom *= n - i + 1
        numer *= i
        div := gcd(denom, numer)
        denom /= div
        numer /= div
    }
    return denom / numer
}

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Split(bufio.ScanWords)
    n, p := nextInt(), nextInt()

    cnt0, cnt1 := 0, 0
    for i := 0; i < n; i++ {
        a := nextInt()
        switch {
        case a%2 == 0:
            cnt0++
        default:
            cnt1++
        }
    }

    ans0 := 0
    ans1 := 0
    for j := 0; j <= cnt0; j++ {
        ans0 += bico(cnt0, j)
    }
    for k := p; k <= cnt1; k += 2 {
        ans1 += bico(cnt1, k)
    }

    fmt.Println(ans0 * ans1)
}

二項係数が最終的に(64bit整数の範囲内の)整数になることを利用して、逐一約分することで有理数の形のまま計算結果を記録する方法です。

誤差の心配がなく、また人力で計算する手順に近いため非常に直感的な解法だと思います。ごり押しとも言いますね。

ただし最小公倍数を毎回求める必要があるため、計算量としてはO(N^3\log N)になると思います。

雑記

  • 想定解法がO(1)なのですべてが霞みます。

  1. \binom{n}{k} = _nC_k(二項係数)