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

Go

入力について

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

高速な入力のためには bufio を使います。bufio を使った方法では 3*106 個の整数の入力にかかる時間は0.1秒以下でした。bufio.Scanner の Scan() と Text() の組み合わせで多くの場合は事足りるのですが，これにも少し問題があります。Scan() で読み取れるトークンの最大長は MaxScanTokenSize という定数で定義されており，Go 1.8 では 64*1024 でした。これでは 105 や 106 程度の長さの文字列の処理を要求する問題に対応することができません。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)
}

偉大な先人の記事

qiita.com

qiita.com

おすすめの本

みんなのGo言語【現場で使える実践テクニック】

みんなのGo言語【現場で使える実践テクニック】

RUPC 2017 Day2 I: Tree-Light

http://judge.u-aizu.ac.jp/onlinejudge/cdescription.jsp?cid=RitsCamp17Day2&pid=I

深さ優先探索をしたときの訪問順序によって頂点を並べると，ある頂点を根とする部分木に含まれる頂点が連続した区間に含まれる性質を使います。以下のページが詳しいです。

これを使えば，木を列として扱うことができるようになるので，列に対して以下の2つのクエリに答えられるようにすればよいです。

  • count(l, r, x, y): [l, r) の中でx以上y以下の値をもつ要素の数を答える
  • change(l, r, x, y): [l, r) の中で値がxである要素をすべてyに書き換える

厄介なのはchangeです。ここで注目するのは，各要素が取り得る値の範囲が0〜9の10段階と小さいことです。ここでchangeを置換として表現してみることを考えましょう。change(l, r, 0, 5) というクエリは (5 1 2 3 4 5 6 7 8 9) と表現できたり，change(l, r, 5, 8) というクエリは (0 1 2 3 4 8 6 7 8 9) と表現できたりします。また，置換は合成することができて，change(l, r, 0, 5) を適用してから change(l, r, 5, 8) を適用するようなクエリは (8 1 2 3 4 8 6 7 8 9) と表現できます。ここで重要なのは，置換という作用素はモノイド（結合法則が成り立ち，単位元が存在する）になっているので遅延セグメント木での作用として使うことができるということです。セグメント木に対して適用可能な操作については次の記事が詳しいです。

koba-e964.hatenablog.com

tomcatowl.github.io

遅延セグメント木と同じ要領で範囲に対する置換を実装した平方分割を以下に示します。これを使うことで O(N+Q√N) で解けます。

typedef array<int, 10> ten;
const int sqrtN = 512;
ten multi(ten a, ten b) {
  ten ret;
  for (int i = 0; i < 10; i++) {
    ret[i] = b[a[i]];
  }
  return ret;
}
struct SqrtDecomposition {
  int N, K;
  vector<int> data;
  // blockCount[k][z]: k番目のブロックの中で値がzである要素の数
  vector<ten> blockCount;
  // 遅延された操作を持っているか？
  vector<bool> lazyFlag;
  // 遅延された操作（置換）
  vector<ten> lazyUpdate;
  // 置換の単位元
  ten initialValue;
  SqrtDecomposition(int n) : N(n) {
    K = (N + sqrtN - 1) / sqrtN;
    data.assign(K * sqrtN, 0);
    lazyFlag.assign(K, false);
    // 置換の単位元は (0 1 2 ... 7 8 9)
    for (int i = 0; i < 10; i++) {
      initialValue[i] = i;
    }
    lazyUpdate.assign(K, initialValue);
    ten zero;
    zero.fill(0);
    zero[0] = sqrtN;
    blockCount.assign(K, zero);
  }
  // k番目のブロックに対する遅延している操作を適用する
  void eval(int k) {
    if(lazyFlag[k]) {
      lazyFlag[k] = false;
      for (int i = k * sqrtN; i < (k + 1) * sqrtN; i++) {
        data[i] = lazyUpdate[k][data[i]];
      }
      lazyUpdate[k] = initialValue;
    }
  }
  // [a, b) の範囲に対して置換を作用させる
  void put(int a, int b, ten perm) {
    for(int k = 0; k < K; ++k) {
      int l = k * sqrtN, r = (k + 1) * sqrtN;
      if(r <= a || b <= l) continue;
      if(a <= l && r <= b) {
        lazyFlag[k] = true;
        // 作用を合成する
        lazyUpdate[k] = multi(lazyUpdate[k], perm);
        // blockCountを更新
        ten cnt;
        cnt.fill(0);
        for(int i = 0; i < 10; i++) {
          cnt[perm[i]] += blockCount[k][i];
        }
        blockCount[k] = cnt;
      }
      else {
        eval(k);
        // 置換を適用
        for(int i = max(a, l); i < min(b, r); ++i) {
          data[i] = perm[data[i]];
        }
        // blockCountを数え直す
        blockCount[k].fill(0);
        for (int i = l; i < r; ++i) {
          blockCount[k][data[i]]++;
        }
      }
    }
  }
  // [a, b) の範囲でx以上y以下の値をもつ要素の数を返す
  int get(int a, int b, int x, int y) {
    int ret = 0;
    for(int k = 0; k < K; ++k) {
      int l = k * sqrtN, r = (k + 1) * sqrtN;
      if(r <= a || b <= l) continue;
      if(a <= l && r <= b) {
        for (int i = x; i <= y; i++) {
          ret += blockCount[k][i];
        }
      }
      else {
        eval(k);
        for(int i = max(a, l); i < min(b, r); ++i) {
          if(x <= data[i] && data[i] <= y) {
            ret++;
          }
        }
      }
    }
    return ret;
  }
};

AtCoder Regular Contest 038 C：茶碗と豆

C: 茶碗と豆 - AtCoder Regular Contest 038 | AtCoder

100点解法

見るからにNimとかGrundy数っぽい雰囲気を感じるので，ググってみると良質な記事がたくさん見つかります。今回は以下の記事を参考にしました。

http://pekempey.hatenablog.com/entry/2016/01/16/190208pekempey.hatenablog.com

最初に，左からi番目の茶碗に豆が1つだけ入っている場合を考えましょう。この場合のGrundy数を grundy[i] という変数に入れておくことにします。i=0 なら動かせる豆が存在しないので負けとなるわけですが，Grundy数の定義より，負けの状態のGrundy数は0なので grundy[0]=0 です。i≠0の場合は豆を i-C[i] 〜 i-1 のどこかに移動させる状態遷移があるので，Grundy数の定義より，grundy[i] は grundy[i-C[i]], grundy[i-C[i]+1], ..., grundy[i-2], grundy[i-1] に含まれない最小の非負整数です。これで，左からi番目の茶碗に豆が1つだけ入っている場合のGrundy数が計算できました。

続いて豆が複数ある場合のGrundy数を計算します。これはGrundy数の魔法により，全ての豆のそれぞれの grundy[i] のXORを取ったものが答えとなります。1つの茶碗に入っている豆の数は最大で109個あるので全ての豆についてGrundy数のXORを取っていると間に合いませんが，同じ数のXORをとると0になるのでmod 2で計算をサボることができます。

Grundy数の計算に O(N2) かかるので計算量は O(N2) です。

104点解法

「grundy[i-C[i]], grundy[i-C[i]+1], ..., grundy[i-2], grundy[i-1] に含まれない最小の非負整数」を計算するところを O(log N) くらいにできれば全体で O(N log N) になるので，この部分の計算を高速化しましょう。

セグメント木の要素にgrundy[i]の値をそのまま持たせて「grundy[i-C[i]], grundy[i-C[i]+1], ..., grundy[i-2], grundy[i-1] に含まれない最小の非負整数」を計算できれば簡単なのですが，残念ながらこれはできそうにありません。

唐突ですが grundy[i] を計算する際に，次のように定義される配列 index を使ってみることを考えましょう。

index[k]：grundy[j]=k となる最大の j (0 <= j < i)

この配列の重要な性質は以下の2つです。

  • 「grundy[i]>=k」 ⇔ 「index[0], index[1], ..., index[k-1] の値が i-C[i] 〜 i-1 の範囲にある」
  • 「index[0], index[1], ..., index[k-1] の値が i-C[i] 〜 i-1 の範囲にある」 ⇔ 「index[0], index[1], ..., index[k-1] の最小値が i-C[i] 以上である」

よって grundy[i]>=k となる最大の k を求めるためには index[0], index[1], ..., index[k-1] の最小値が i-C[i] 以上となる最大の k を求めれば良いことが分かります。これをナイーブに実装すると次のようになります。

grundy := make([]int, N)
grundy[0] = 0
index := make([]int, N)
index[0] = 0
for i := 1; i < N; i++ {
    index[i] = -1
}
for i := 1; i < N; i++ {
    for k := 0; k < N; k++ {
        if i - C[i] <= index[k] {
            // do nothing
        } else {
            grundy[i] = k
            break
        }
    }
    index[grundy[i]] = i
}

これでは O(N2) になってしまうので「index[0], index[1], ..., index[k-1] の最小値が i-C[i] 以上となる最大の k」の計算をセグメント木で実装します。私は平方分割で実装しました。次の記事も参考にしてみてください。Flipping Parenthesesの説明が参考になるかもしれません。

kujira16.hateblo.jp

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

const (
    sqrtN = 512
    INF = 1000000000
)

type SqrtDecomposition struct {
    N        int
    K        int
    data     []int
    blockMin []int
}

func NewSqrtDecomposition(n, initialValue int) *SqrtDecomposition {
    k := (n + sqrtN - 1) / sqrtN
    sq := &SqrtDecomposition{
        N:        n,
        K:        k,
        data:     make([]int, k*sqrtN),
        blockMin: make([]int, k),
    }
    for i := 0; i < len(sq.data); i++ {
        sq.data[i] = initialValue
    }
    for i := 0; i < len(sq.blockMin); i++ {
        sq.blockMin[i] = initialValue
    }
    return sq
}

// data[x]=y
func (sq *SqrtDecomposition) Update(x, y int) {
    k := x / sqrtN
    sq.data[x] = y
    minval := INF
    for i := k * sqrtN; i < (k+1)*sqrtN; i++ {
        minval = intMin(minval, sq.data[i])
    }
    sq.blockMin[k] = minval
}

// min [0, k) >= x となる最大のkを返す
func (sq *SqrtDecomposition) Find(x int) int {
    for k := 0; k < sq.K; k++ {
        if sq.blockMin[k] < x {
            for i := k * sqrtN; i < (k+1)*sqrtN; i++ {
                if sq.data[i] < x {
                    return i
                }
            }
        }
    }
    return -1
}

これを使うと先ほどのナイーブな O(N2) の方法が以下のように書き直せます。

grundy := make([]int, N)
grundy[0] = 0
sq := NewSqrtDecomposition(N, -1)
sq.Update(0, 0)
for i := 1; i < N; i++ {
    grundy[i] = sq.Find(i - C[i])
    sq.Update(grundy[i], i)
}

セグメント木を使えば計算量は O(N log N) です。今回は平方分割なので O(N√N) です。

Submission #1186675 - AtCoder Regular Contest 038 | AtCoder

参考にした記事

mayokoex.hatenablog.com