AtCoderで使えるC++STLライクなライブラリ「GoSTL」を紹介

Aru

AtCoderのGo環境にはいくつかのライブラリがインストールされています。その中の１つにC++のSTLライクな機能を実現するGoSTLがあります。この記事では、GoSTLについて簡単に説明します

GoSTLとは

GoSTLはGithubに公開されているGo言語のライブラリです。

GoSTL is a data structure and algorithm library for go, designed to provide functions similar to C++ STL, but more powerful. Combined with the characteristics of go language, most of the data structures have realized goroutine-safe. When creating objects, you can specify whether to turn it on or not through configuration parameters.

（以下、翻訳）
GoSTL は Go 用のデータ構造およびアルゴリズムライブラリであり、C++ STL に似ていますが、より強力な機能を提供するように設計されています。 go 言語の特性と相まって、ほとんどのデータ構造は goroutine-safe を実現しています。オブジェクトを作成するときに、構成パラメータを使用してオンにするかどうかを指定できます。

https://github.com/liyue201/gostl

ソースコードとドキュメントは以下のリンクにあります。

Githubリポジトリ：https://github.com/liyue201/gostl
ドキュメント：https://pkg.go.dev/github.com/liyue201/gostl

ドキュメントは、全部は揃っていないようなので今後整備されることに期待です。

眺めてみるとC++STLのライブラリの主要なものは揃っているようです。Go言語を使っていて欲しいのは、sortedSet、multisetあたりだと思いますが、これも揃っています。

この記事では、setのLowerBoundを使った例をAtCoder ABC-370C問題を解いてみます。

ローカル環境でライブラリを使う場合の注意点

外部のライブラリを使う場合、go.modを作成する必要があります。いろいろやり方はありますが以下の手順が楽ですので紹介しておきます。

go mod init xxxx（xxxxは任意）で、go.modを生成する
ソースコードにライブラリを追加する
"github.com/liyue201/gostl/ds/set"など
※VSCodeなどを利用している場合、コード中でライブラリの関数を使っていないとimportが削除されることに注意
go mod tidyで、go.modを更新

以上の手順で、パッケージが利用できるようになります。

ABC-370D問題を解く

GoSTLには色々なアルゴリズムが入っていますが、今回はSetの使い方を紹介します。具体的には、GoSTLのSetを使ってAtCoderのABC370-D問題を解いてみました。

どんな問題？

爆弾を置いて壁を爆破する問題です。最初は全てのマスに壁があって、壁のあるマスに爆弾を設置するとマスの壁が消えます。壁がないマスに爆弾を設置すると、上下左右の一番近い壁が破壊されるというものです。

ある行だけで考えた場合、「爆弾を置いた場所から一番近い壁」を高速に検索する問題になります。

壁がある場所を管理しておいて、爆弾を置いた場所に近い壁を二分探索すればOKです。ただ、壁が破壊されると削除する必要があるため、配列で管理すると削除コストが大きいためTLEします。

C++であればSTLのsetを使えばよいですがGo言語にはSortedSetがないので実装は大変です。ただ、Go言語もGoSTLを使えば同じように簡単に解くことができます。

setオブジェクトの生成

ジェネリクス対応になっているため、慣れないとわかりにくいかもしれません

追加（insert）, 削除（erase）

追加と削除はSTLに近いです。Go言語なので、先頭の文字が大文字のInsertとEraseになることに注意です。

値の参照（Value）

C++と違って、*itという形で参照できないので値を参照する場合はValue()メソッドを使う必要があります。C++STLでもvalue()を使うことがあるので違和感はなかったです。

イテレータを進める（next)、戻る(prev)

C++では、it++などで移動させることができますが、GoSTLではprev(), next()を使います。 it.Prev()のように戻り値なしでOKのようです。

説明があまりないので、コードをみながら使い方を調べる必要があるかもしれません

プログラム全体は以下になります。

package main

import (
	"bufio"
	"fmt"
	"math"
	"os"
	"strconv"

	"github.com/liyue201/gostl/ds/set"
	"github.com/liyue201/gostl/utils/comparator"
)

var sc = bufio.NewScanner(os.Stdin)
var wr = bufio.NewWriter(os.Stdout)

func out(x ...interface{}) {
	fmt.Fprintln(wr, x...)
}

func getI() int {
	sc.Scan()
	i, e := strconv.Atoi(sc.Text())
	if e != nil {
		panic(e)
	}
	return i
}

func main() {
	defer wr.Flush()
	sc.Split(bufio.ScanWords)
	sc.Buffer([]byte{}, math.MaxInt32)
	// this template is new version.
	// use getI(), getS(), getInts(), getF()
	H, W, Q := getI(), getI(), getI()

	h := make([]*set.Set[int], H)
	for i := 0; i < H; i++ {
		h[i] = set.New[int](comparator.IntComparator, set.WithGoroutineSafe())
		h[i].Insert(-1)
		h[i].Insert(W)
		for j := 0; j < W; j++ {
			h[i].Insert(j)
		}
	}
	w := make([]*set.Set[int], W)
	for i := 0; i < W; i++ {
		w[i] = set.New[int](comparator.IntComparator, set.WithGoroutineSafe())
		w[i].Insert(-1)
		w[i].Insert(H)
		for j := 0; j < H; j++ {
			w[i].Insert(j)
		}
	}

	erase := func(i, j int) {
		h[i].Erase(j)
		w[j].Erase(i)
	}

	// m := make(map[pos]bool)
	for qi := 0; qi < Q; qi++ {
		r, c := getI()-1, getI()-1

		{
			it := h[r].LowerBound(c)
			if it.Value() == c {
				erase(r, c)
				continue
			} else {
				if it.Value() != W {
					erase(r, it.Value())
				}
				it = h[r].LowerBound(c)
				it.Prev()
				if it.Value() != -1 {
					erase(r, it.Value())
				}
			}
		}
		{
			it := w[c].LowerBound(r)
			if it.Value() != H {
				erase(it.Value(), c)
			}
			it = w[c].LowerBound(r)
			it.Prev()
			if it.Value() != -1 {
				erase(it.Value(), c)
			}
		}
	}

	ans := 0
	for i := 0; i < H; i++ {
		ans += h[i].Size() - 2
	}
	out(ans)
}