bufio で string を取ろうとして詰まった話

RE?

Atcoder で https://atcoder.jp/contests/abc237/tasks/abc237_d を解いていて、
テストケースのうち数ケースのみ RE ( Runtime Error ) になってしまった。

原因

自分の実装

bufio の scanner をこんな感じで使っていたのが原因

var sc = bufio.NewScanner(os.Stdin)

func main(){
    sc.Split(bufio.ScanWords)
    s:=scanString()
}

func scanString() (s string) {
	sc.Scan()
	s = sc.Text()
	return
}

ソースコード

定数

const (
	// MaxScanTokenSize is the maximum size used to buffer a token
	// unless the user provides an explicit buffer with [Scanner.Buffer].
	// The actual maximum token size may be smaller as the buffer
	// may need to include, for instance, a newline.
	MaxScanTokenSize = 64 * 1024

	startBufSize = 4096 // Size of initial allocation for buffer.
)

MaxScanTokenSize を最大のバッファサイズ
startBufSize をはじめに確保されるバッファのサイズ
として定義されている

bufio.NewScanner の実装

func NewScanner(r io.Reader) *Scanner {
	return &Scanner{
		r:            r,
		split:        ScanLines,
		maxTokenSize: MaxScanTokenSize,
	}
}

bufio.NewScanner を呼ぶと、Scanner に対して maxTokenSize が設定される

func (s *Scanner) Scan() bool の実装

func (s *Scanner) Scan() bool {
    
    // 省略
    ...
    
    // バッファサイズが最大トークンサイズを超えていないかチェック
    if len(s.buf) >= s.maxTokenSize || len(s.buf) > maxInt/2 {
        s.setErr(ErrTooLong)
        return false
    }
    
    ...
    // 省略
    ...
    
    // バッファサイズの初期化
    if newSize == 0 {
        newSize = startBufSize
    }
    
    ...
    // 省略
}

上のチェックの部分は、今回のエラーの原因であると考えられる
下の部分で、バッファが確保されていない場合 ( 初期の状態 ) 、startBufSize の 4096 bytes が指定される

解決

今回の Atcoder の問題では、入力される string のサイズの制約が

1\le n \le5×10^5

となっていたので デフォルトのバッファサイズの最大値である 64 * 1024 = 65,536 を超えてしまった
以下のコードを Scan() を行う前に記述することで解決した

sc.Buffer(make([]byte, 64*1024), 1e6)

これは、scanner が使用するバッファを指定するもの

• 第 1 引数の []byte は初期バッファ
  今回は、デフォルトの最大 64 * 1024 = 65,536 を超えることがわかっていたので、初めから 64 * 1024 のサイズを確保した

• 第 2 引数の 1e6 は最大バッファサイズ
  最大 1*10^6 文字と少し多めにとっている
  開発に使用するなら、2 の累乗単位で取るべきなのかもしれないが、今回は競技プログラミングなので深く考えなかった

まとめ

競技プログラミングで bufio パッケージを使用するときは、適切なバッファサイズの設定が重要。
特に大量のデータを扱う場合には、事前に適切なサイズを確保する。