fmt.ScanやbufioのScannerで標準入力を受け取る

競プロで標準入力を受け取るとき、行単位で受け取ったりスペース区切りで受け取ったりしたいけど、どうすれば良いんだっけってなったのでまとめておく。

fmt.Scan

結論、fmt.Scanで全て解決する。

Scan scans text read from standard input, storing successive space-separated values into successive arguments. Newlines count as space. It returns the number of items successfully scanned. If that is less than the number of arguments, err will report why.

pkg.go.dev

スペース区切りで変数に格納するし、改行もスペースとして扱うらしい。

fmt.Scanの引数はポインタかScannerインターフェースでなければいけないらしいので注意。

All arguments to be scanned must be either pointers to basic types or implementations of the Scanner interface.

pkg.go.dev

実行例はこんな感じ。

var s1, s2, s3 string

fmt.Scan(&s1, &s2, &s3)
fmt.Printf("Your input is %s %s %s\n", s1, s2, s3)

❯ go run cmd/main.go
1 2 3
Your input is 1 2 3

❯ go run cmd/main.go
1
2
3
Your input is 1 2 3

可変長で標準入力を受けたければ、以下のような感じ。

var N int
fmt.Scan(&N)
a := make([]int, N)
for i := 0; i < N; i++ {
    fmt.Scan(&a[i])
}
fmt.Println(a)

❯ go run tmp/main.go
3
1 2 3
[1 2 3]

fmt.Scanで困るパターン

便利なfmt.Scanだけど実は遅くて困る時がある。

実際にAtCoder Beginner Contest 246の C - Couponでは、1行で空白区切りの文字を最大で10⁹個読み込む必要がある。

以下の画像のようにfmt.Scanで表順入力を受け取るかその他の高速な方法で受け取るかで実行速度や使用するメモリの量が変わる。

fmt.ScanとbufioのScannerを使った場合の実行速度や使用メモリの違い

以下の問題でもfmt.Scanではタイムアウトしてしまう。 atcoder.jp

bufio.Scanner

そこで便利なのがbufioパッケージのScanner。

使い方は以下。初めに整数Nを受け取り、N個の空白区切りの整数を受け取るとする。

// ひとまず初期化
var scanner = bufio.NewScanner(os.Stdin)

func nextInt() int {
    scanner.Scan() // トークンの読み込み
    i, e := strconv.Atoi(scanner.Text()) // 読み込んだトークンをTextメソッドで取り出し
    if e != nil {
        panic(e)
    }
    return i
}

func main() {
    var N int
    fmt.Scan(&N)
    scanner.Split(bufio.ScanWords) // 空白区切りで受け取るので引数にbufio.ScanwWordsを渡す
    A := make([]int, N)
    for i := 0; i < N; i++ {
        A[i] = nextInt() // ScanメソッドやTextメソッドで読み込んで取り出して加工
    }