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

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

AtCoder AGC 002 B - Box and Ball (Go)

irisruneです。久々の過去問回です。

問題

atcoder.jp

突き詰めればただシミュレートするだけの問題で、操作によってどのような挙動を示すのかを把握できれば解くことができると思います。

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
}

type box struct {
    balls int
    red   int
}

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

    boxes := make([]box, n+1)
    for i := range boxes {
        boxes[i].balls = 1
    }
    boxes[1].red = 1

    for j := 0; j < m; j++ {
        x, y := nextInt(), nextInt()
        boxes[y].balls++
        if boxes[x].red == 1 {
            boxes[y].red = 1
        }
        boxes[x].balls--
        if boxes[x].balls == 0 {
            boxes[x].red = 0
        }
    }

    ans := 0
    for _, b := range boxes[1:] {
        ans += b.red
    }
    fmt.Println(ans)
}

各箱についてボールが何個あるか、赤いボールが入っている可能性があるかを管理しておきます。この状態で操作を行ったとき、以下のことがいえます。

  1. ボールを入れた箱のボールの数は1個増える。
  2. ボールを取り出した箱に赤いボールが入っている可能性があった場合、ボールを入れた箱のボールに赤いボールが入っている可能性がある。
  3. ボールを取り出した箱のボールの数は1個減る。
  4. ボールを取り出した箱のボールの数が0個になる場合、ボールを取り出した箱に赤いボールが入っている可能性はない。

以上を踏まえて、各箱のボールの数を1個、箱1のみ赤いボールが入っている可能性があるという初期状態から操作をシミュレートすることで解くことができます。ただし、上記2.の処理の前に4.の処理を行うなどすると結果が狂う可能性があるため気を付けましょう。

雑記

  • 前回の問題(ABC135D)ですが無事ACできました。
    • 前にも引っかかった入力方法による文字列長の上限によって、長大文字列を途中で打ち切って入力していたのが原因です。