Go言語で競技プログラミングするときに使ってる入出力

入力について

普通は fmt.Scan で問題ないのですが，高速な入力が要求される問題では fmt.Scan だと遅すぎることがあります。fmt.Scan では 3*10⁶ 個の整数の入力に10秒以上かかりました。

高速な入力のためには bufio を使います。bufio を使った方法では 3*10⁶ 個の整数の入力にかかる時間は0.1秒以下でした。bufio.Scanner の Scan() と Text() の組み合わせで多くの場合は事足りるのですが，これにも少し問題があります。Scan() で読み取れるトークンの最大長は MaxScanTokenSize という定数で定義されており，Go 1.8 では 64*1024 でした。これでは 10⁵ や 10⁶ 程度の長さの文字列の処理を要求する問題に対応することができません。Buffer([]byte, int) でバッファの大きさを変えることができますが，バッファの長さを変えるのを忘れてWAになってしまいそうです。

そこで bufio.Reader の ReadLine() を使います。戻り値の isPrefix を確認して ReadLine() を複数回呼び出すことで長い文字列でも読み取ることができます。

出力について

os.Stdout はバッファされないため，毎回flushしているのと同じような状態で遅いです。os.Stdout を直接利用している fmt.Printf なども同様です。これを解決するために bufio を利用しています。プログラム終了前にはflushしてやる必要がありますが，忘れやすいのでmainの先頭で defer io.Flush() としてやるほうが良いと思います。

配列の中身を空白区切りで出力したい場合が多々ありますが，...interface{} の実引数として []int や []string を与えることはできません。そのため []intと[]stringのためだけに配列の中身を空白区切りで出力する関数を定義しています。

[]int や []string から []interface{} への変換についての話題が気になる方は以下のページを参照してください。

InterfaceSlice · golang/go Wiki · GitHub

その他

maxやminの関数はint用だけ特別に用意しています。Go言語にはジェネリクスが導入される予定がないので筋肉で解決しましょう。以下のリポジトリがすごくおもしろいです。

github.com

Go言語のPrintfのフォーマット指定子には，構造体の中身が見られる ‘%+v’ というものがあり，デバッグに便利です。printfデバッグするためにLog関数を定義しています。IntelliJ IDEAのデバッガが使いやすいのでprintfデバッグを使わないことも多いですが。

ソースコード

package main

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

type Io struct {
    reader    *bufio.Reader
    writer    *bufio.Writer
    tokens    []string
    nextToken int
}

func NewIo() *Io {
    return &Io{
        reader: bufio.NewReader(os.Stdin),
        writer: bufio.NewWriter(os.Stdout),
    }
}

func (io *Io) Flush() {
    err := io.writer.Flush()
    if err != nil {
        panic(err)
    }
}

func (io *Io) NextLine() string {
    var buffer []byte
    for {
        line, isPrefix, err := io.reader.ReadLine()
        if err != nil {
            panic(err)
        }
        buffer = append(buffer, line...)
        if !isPrefix {
            break
        }
    }
    return string(buffer)
}

func (io *Io) Next() string {
    for io.nextToken >= len(io.tokens) {
        line := io.NextLine()
        io.tokens = strings.Fields(line)
        io.nextToken = 0
    }
    r := io.tokens[io.nextToken]
    io.nextToken++
    return r
}

func (io *Io) NextInt() int {
    i, err := strconv.Atoi(io.Next())
    if err != nil {
        panic(err)
    }
    return i
}

func (io *Io) NextFloat() float64 {
    i, err := strconv.ParseFloat(io.Next(), 64)
    if err != nil {
        panic(err)
    }
    return i
}

func (io *Io) PrintLn(a ...interface{}) {
    fmt.Fprintln(io.writer, a...)
}

func (io *Io) Printf(format string, a ...interface{}) {
    fmt.Fprintf(io.writer, format, a...)
}

func (io *Io) PrintIntLn(a []int) {
    b := []interface{}{}
    for _, x := range a {
        b = append(b, x)
    }
    io.PrintLn(b...)
}

func (io *Io) PrintStringLn(a []string) {
    b := []interface{}{}
    for _, x := range a {
        b = append(b, x)
    }
    io.PrintLn(b...)
}

func Log(name string, value interface{}) {
    fmt.Fprintf(os.Stderr, "%s=%+v\n", name, value)
}

func intMin(a, b int) int {
    if a < b {
        return a
    }
    return b
}

func intMax(a, b int) int {
    if a > b {
        return a
    }
    return b
}

func main() {
    io := NewIo()
    defer io.Flush()
    A := io.NextInt()
    B := io.NextInt()
    C := io.NextInt()
    S := io.Next()
    io.PrintLn(A+B+C, S)
}