【競プロ】ScanningでTLEが発生してコケる

はじめに

Go言語で競技プログラミングに臨んだことがある人なら, 一度はfmt.Scan関数の落とし穴にハマったことがあるのではないでしょうか. Go言語では標準入力を受け付ける際に, fmt.Scan関数, もしくはbufio.Scannerメソッドやbufio.ReadLineメソッドをよく用いります.
しかし, Go言語を使って競技プログラミングをやる場合, bufio.Scannerを使用しなければ困る場面があります.
今回の記事では, fmtパッケージのScan関数とbufioパッケージのScannerメソッドに注目しながら, fmt.ScanでTLEが発生する原因と解決策について紹介します.

結論だけ知りたい人はこちらをクリック

fmt.Scan

まずは, fmt.Scan関数の詳細を見てみます.

fmt.ScanでTLEが発生する

競プロにおいて, 大量のデータ入力が必要になる場合があります.
例えば, 競プロでよく見かける以下の入力形式において

N
A1 A2 A3 ... AN

ここで, レギュレーションとして N の最大値を $3*10^6$ とします.
一般的にGo言語ではこのような形式で入力を受け付ける時,

func main() {
	N := // Nの入力を受け付ける・・・①
	if N > int(3*(math.Pow(10, 6))) {
		log.Fatalf("Regulation error: N => %d\n", N)
		return
	}

	A_list := make([]int, N)
	for i := 0; i < N; i++ {
		A_list[i] = // Aの入力を受け付ける・・・②
	}
}

のような書き方をするのではないでしょうか.
この時, N の最大値は $3*10^6$ なので, for文の中では300万件ものデータを受け付けなければなりません.
そのため, このような状況において②の箇所でfmt.Scanを使うと多くの場合, TLE（Time Limit Exceeded）による不正解判定が出てしまいます.

なぜ fmt.Scan だとダメなのか

なぜ, fmt.ScanだとTLEするのかを理解するために, fmtパッケージのscan.goを読みに行ってみます.

// 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.
//
func Scan(a ...any) (n int, err error) {
	return Fscan(os.Stdin, a...)
}

Scanは, 標準入力から読み込んだテキストをスキャンし, スペースで区切られた値として, 順に引数に格納します. 改行文字はスペースとしてカウントされます. この関数はスキャンに成功した項目数を返します. この数が, 引数の数より少ないときは, errにその理由を返します.

// Fscan scans text read from r, 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.
//
func Fscan(r io.Reader, a ...any) (n int, err error) {
	s, old := newScanState(r, true, false)
	n, err = s.doScan(a)
	s.free(old)
	return
}

Fscanは, r (io.Reader) から読み取られたテキストをスキャンし, スペースで区切られた連続する値を連続する引数に格納します. 改行はスペースとしてカウントされます. 正常にスキャンされたアイテムの数を返します. それが引数の数より少ない場合, errはその理由を報告します.

fmt.Scanは内部的にfmt.Fscan関数を呼び出します. つまり, fmt.Scanはfmt.Fscanのラッパー関数であり, 引数のio.Readerに標準入力を渡しているだけです. I/O処理はプログラムを実行する上では比較的重い動作であり, unbufferd I/Oの場合, 内部バッファを使用しずに直接ディスクに対してファイル書き込みを行うため, 都度システムコールが発生します. そのため, $3 * 10^6$ 程度の数値を入力するとTLEします.
つまり, 通常は比較的簡単に書くことができるfmt.Scanで問題なくても, 競プロのように膨大な量のデータを高速で受け付ける必要がある場合は, fmt.Scanだと遅すぎる訳です.

bufio.Scanner

では, bufio.Scannerの場合はどうなっているのか.
まず, Scannerの正体を知るために, bufioパッケージのScanner構造体を見にいきます.

// Scanner provides a convenient interface for reading data such as
// a file of newline-delimited lines of text. Successive calls to
// the Scan method will step through the 'tokens' of a file, skipping
// the bytes between the tokens. The specification of a token is
// defined by a split function of type SplitFunc; the default split
// function breaks the input into lines with line termination stripped. Split
// functions are defined in this package for scanning a file into
// lines, bytes, UTF-8-encoded runes, and space-delimited words. The
// client may instead provide a custom split function.
// 
type Scanner struct {
	r            io.Reader // The reader provided by the client.
	split        SplitFunc // The function to split the tokens.
	maxTokenSize int       // Maximum size of a token; modified by tests.
	token        []byte    // Last token returned by split.
	buf          []byte    // Buffer used as argument to split.
	// 以下, 省略
}

Scannerは, 改行で区切られたテキスト行のファイルなどのデータを読み取るための便利なインターフェイスを提供します. Scanメソッドを連続して呼び出すと, ファイルの「トークン」がステップスルーされ, トークン間のバイトがスキップされます. トークンの指定は, SplitFunc型のsplit関数によって定義されます. デフォルトの分割関数は, 入力を行の終端を取り除いた行に分割します. このパッケージでは, ファイルを行, バイト, UTF-8でエンコードされたルーン文字, およびスペースで区切られた単語にスキャンするための分割関数が定義されています. クライアントは, 代わりにカスタム分割機能を提供する場合があります.

Scannerメソッドは, トークン によって, 単語ごと（スペース区切り）, 行ごと（改行文字区切り）に効率的に読み取ることができます. また, トークンの指定はSplitFunc型のsplit関数によって定義することができるようです.
split関数（無名関数）は以下のような定義になっています.

// SplitFunc is the signature of the split function used to tokenize the input. 
type SplitFunc func(data []byte, atEOF bool) (advance int, token []byte, err error)

SplitFuncは, 入力をトークン化するために使用される分割関数のシグネチャです.

SplitFunc型に代入することができる関数として, 以下の4つが定義されています.

// バイト ごと
func ScanBytes(data []byte, atEOF bool) (advance int, token []byte, err error)

// 行 ごと
func ScanLines(data []byte, atEOF bool) (advance int, token []byte, err error)

// ルーン（int32のエイリアス）ごと
func ScanRunes(data []byte, atEOF bool) (advance int, token []byte, err error)

// 単語 ごと
func ScanWords(data []byte, atEOF bool) (advance int, token []byte, err error)

これらの関数は, split関数とunderlying typeが同一であるため, Go言語の特性の一つである Assignability によって代入することが可能になっています.

Assignabilityに関しては, こちらの記事が分かりやすいかと思います.

加えて, Scannerメソッドはfmt.Scanと異なり, 内部にバッファを持っています. これを利用することで, splitフィールドに定義されたトークンに従ってバッファに格納されたバイト列ごとにI/O処理（bufferd I/O）を行うため, fmt.Scanよりも高速に動作することが可能になります.

スキャナーの作成は, NewScanner関数によって行います.

// NewScanner returns a new Scanner to read from r.
// The split function defaults to ScanLines.
func NewScanner(r io.Reader) *Scanner

デフォルトで作成されたスキャナーは, 「行」をトークンとして, 入力を受け付けるよう設定されています.
これを上述したsplit関数に置き換えたい場合, Splitメソッドを使用します.

// Split sets the split function for the Scanner.
// The default split function is ScanLines.
func (s *Scanner) Split(split SplitFunc)

実際に使用する場合

これを以下のような形式で受け付ける場合,

N
A1 A2 A3 ... AN

単語（空白スペース）ごとに受け付けることで, 高効率かつ高速に処理が可能になります.

// スキャナを作成
var sc = bufio.NewScanner(os.Stdin)

func main() {
	sc.Split(bufio.ScanWords) // 単語（空白スペース）で分割
    for i := 0; i < N; i++ {
		A_list[i] = scanInt() // bufio.Scannerを呼び出す
	}
}    

// scanInt scans int type.
func scanInt() (i int) {
	var err error
	sc.Scan()
	i, err = strconv.Atoi(sc.Text()) // int型にparse
	if err != nil {
		log.Fatal("Failed to input", err)
	}
	return
}

速度比較

実際に以下のコードを使って, データの入力値を増やした時の, fmt.Scanとbufio.Scannerの読み込み速度を比較してみます. （本番環境を想定して, メモリを512MBに制限して計測を行なっています. ）

• https://github.com/GotoRen/samples/tree/master/scanning-speed

データ数	fmt.Scan	bufio.Scanner
10	0.000123s	0.000023s
100	0.000310s	0.000025s
1000	0.004084s	0.000224s
10000 (=$10^4$)	0.033157s	0.000550s
100000 (=$10^5$)	0.509994s	0.006163s
1000000 (=$10^6$)	3.771113s	0.064092s
2000000 (=$2*10^6$)	7.210549s	0.142498s
3000000 (=$3*10^6$)	10.357201s	0.190407s
4000000 (=$4*10^6$)	13.958481s	0.253358s

データ件数が1万（$10^4$）件ぐらいまでは, 双方でほぼ同じ読み込み速度であるため, 簡単に書けるfmt.Scanの方が良さそうです.
しかし, 10万（$10^５$）件を超えたあたりからfmt.Scanが明らかに遅くなってきます. 今回, レギュレーションとして設定されている『メモリ上限：512MB / 実行タイムアウト：2000ms』において, 100万（$10^6$）件を読み込む時点で3秒以上も掛かるようでは, 300万（$3 * 10^6$）件を読み込んで処理する際にTLEが発生するのも納得です.
一方, bufio.Scanはデータ件数が増えた場合にも, 読み込み速度が劣化しにくく安定して使えるため, 膨大な量のデータを読み込むことが可能になります.

結論

これらのことから, ①では, fmt.Scanによって入力を受け付ければ良く, ②では, bufio.Scannerを使用する必要があるということになります.

// ① Nの入力を受け付ける場合
func main() {
	var N int
	fmt.Scan(&N) // calling fmt.Scan
}

// ② 3*10^6回受け付ける場合
var sc = bufio.NewScanner(os.Stdin)

func main() {	
    sc.Split(bufio.ScanWords) // split by word
	for i := 0; i < N; i++ {
		A_list[i] = scanInt() // calling bufio.Scanner
	}
}

// scanInt scans int type.
func scanInt() (i int) {
	var err error
	sc.Scan()
	i, err = strconv.Atoi(sc.Text()) // parse to int
	if err != nil {
		log.Fatal("Failed to input", err)
	}
	return
}

まとめると,

---

• 速度を気にしなくても良い場合
  • fmt.Scan
• 膨大なデータ（$> 10^{5}$）を高速に読み込む必要がある場合
  • bufio.Scanner
• また, 今回は触れなかったが, 行数が多い場合（$> 64*10^{3}$）
  • bufio.ReadLine

---

という感じでしょうか.

うーん. 単純な入力受付一つとっても, プログラミングって奥が深い...
競プロハマりそう. 研究もせねば...

参考・引用

• go/src/fmt/scan.go
• pkg.go.dev/fmt/#Scan
• go/src/bufio/scan.go
• pkg.go.dev/bufio/#Scanner
• toVersus/get_max_profit.go
• C言語におけるファイル入出力の高速化
• Golangのバッファってよくできてるよな
• Go言語のベンチマークでパフォーマンス測定
• Go言語で標準入力から読み込む競プロのアレ
• Goの代入可能性についてちゃんと理解してみる

読者の皆さんへ

本記事を読んで下さり, ありがとうございます.
誤字脱字や誤った内容を見つけた場合, 本記事にてフィードバックを下さると幸いです.
また, 意見や提案等がありましたら, 併せて本記事のコメント欄に記述して頂けると助かります. より良い記事を書くため, 改良の参考にさせて頂きます.

@ren510dev