この記事はand factory.inc Advent Calendar 2023の6日目の記事です。
昨日は @MatsuNaoPen さんの「使って便利なmermaid記法(flowchart編）」でした。
急遽埋めたので完全に業務に関係ない話で恐縮です…。もう一枠取っているのでそちらではもう少し業務絡みの事を書ければと思います。
概要
この記事ではGoで競技プログラミングをする時に作っておくよい関数をお伝えしたいと思います。
Goで競技プログラミングをするに当たってのテクニックは既に詳しくまとめていただいている方がおり、こちらがとても参考になるかと思います。
上記を踏まえた上で実際にどんな関数があると便利かは日々試行錯誤しているのですが、その中でもよく使うお気に入りの関数や作っておいて良かったと思う関数が多々あります。
AtCoderで特にABCに出ている方なら共感していただけるかと思うのですが、ともかく競技中時間と余裕がないです。落ち着いていれば書ける普段の業務ならわざわざ関数化しない処理でも、関数化して考えずにスッと使えるようにしておくとそれが生きる局面は多いように思います。この記事を参考に、是非色々な関数化を試みていただければ幸いです。
登場するすべての関数は以下で都度アップデートしているので、使えそうな関数があれば参考にしつつ是非使ってみていただければと思います。
https://github.com/gosagawa/atcoder/blob/master/_template/main.go
上記のテンプレート含めてどんな感じで普段やっているかは別途まとめているので、まずGoでAtCoderをやってみたいという方は以下も参考にしていただけると思います。
著者のレベル
現状AtCoderだと青を彷徨っています。Goメインで黄色の人はいなさそうなのでそこを目指しています。黄色までは意地でもGoで頑張ろうとしているのですが、なれたら憑物が落ちたようにC++を使い出すかもしれません…。
AtCoder Problemsによると現状だと4番目にGoでACしているようです。なのでライブラリではそれなりにGoで困ることは潰されているかと思います。
入力
どんな問題でもほぼ確実にあるのが入力です。入力を取り込むためのロジックは基本複雑でないですが、出てくるパターンが限られているのでパターン化して関数を用意しおくことは非常に強力です。
基本的な実装と全関数
以下のようにしてバッファからinputを取ります。これをベースに色々な入力をさっと取れるように関数を整えています。早く打ち込めるように関数名はniにしています。最初にどなたかのを参考にして作ったのですがおそらくnew intなのではと思っています。
var sc = bufio.NewScanner(os.Stdin)
func init() {
sc.Buffer([]byte{}, math.MaxInt64)
sc.Split(bufio.ScanWords)
}
func ni() int {
sc.Scan()
return atoi(sc.Text())
}
全部列挙するとこれだけ関数を作っていました。
自分用ライブラリのio関連部分
// intを取り込み
func ni() int
// intを二つ取り込み
func ni2() (int, int)
// intを三つ取り込み
func ni3() (int, int, int)
// intを四つ取り込み
func ni4() (int, int, int, int)
// intのスライスを取り込み
func nis(arg ...int) []int
// intのスライスを二つ取り込み
func ni2s(n int) ([]int, []int)
// intのスライスを三つ取り込み
func ni3s(n int) ([]int, []int, []int)
// intのスライスを四つ取り込み
func ni4s(n int) ([]int, []int, []int, []int)
// 長さ2の配列(array)のスライスを取り込み
func ni2a(n int) [][2]int
// 長さ3の配列(array)のスライスを取り込み
func ni3a(n int) [][3]int
// 長さ4の配列(array)のスライスを取り込み
func ni4a(n int) [][4]int
// 二次元（2 dimention）のスライスを取り込み
func ni2d(n, m int) [][]int
// float64の取り込み
func nf() float64
// stringの取り込み
func ns() string
// stringをintのスライスにして取り込み
func nsis() []int
// 複数のstringを二次元のintのスライスにして取り込み
func nsi2s(n int) [][]int
// nsi2sで取得したものに対して、特定の文字が特定のintになるように変える
// mp = convidxi2s(nsi2s(n), map[string]int{".": 0, "#": 1})
func convidxi2s(sl [][]int, conv map[string]int) [][]int
// intのスライスに対して、特定の文字が特定のintになるように変える
func convidxis(sl []int, conv map[string]int) []int
// 一文字のstringをintにかえる
func ctoi(c string) int
// 数字の文字列をintのスライスに変える
// 20桁以上のintで取り込むとオーバーフローするものとかで利用する
func nsiis() []int
// floatのものを10^n倍してintにしたものを取り込む
func nftoi(decimalLen int) int
// 標準入力からintを読み込む
func scani() int
// 標準入力からstringを読み込む
func scans() string
要素数+数列
このように要素数と数列が入力されるものはよくありますが、スライスを取れる関数を用意しておりさっと二行で取得できます。
5
1 2 3 4 5
↓
n := ni() // new int
as := nis(n) // new int slice
要素数+二つ以上数列
二要素以上がわたることも頻出です
3
1 2
3 4
5 6
以下のようにして、intのスライスを二つ作るか
n := ni()
as, bs := ni2s(n) // new int two slice as: []int{1,3,5} bs: []int{2,4,6}
もしくは以下のようにして[2]intのスライスに投入できます
n := ni()
abs := ni2a(n) // new int two array abs : [][2]int{[2]int{1,2},[2]int{3,4},[2]int{5,6}}
二次元スライス
このように縦横の二次元で渡されるケースもさっと[][]intを作れるようにしてます
3 4
1 2 3 4
5 6 7 8
9 10 11 12
↓
h, w := ni2()
mp := ni2d(h,w) // new int two dimension slice
文字列
文字列が渡されるパターンも出題されます。
abcde
stringとして取得できます。
s := ns() // s: abcde
一方intのスライスに変換するとも多いです。こうすることで扱いやすくなり便利な局面が多いです。例えば、abcdeは[]int{0,1,2,3,4}になります。
as := nsis() // as: []int{0,1,2,3,4}
二次元文字列
このように.と#で表現された平面図（迷路とか)を取り込む問題もたまに出ます。以下だとSからGまでの最短経路は？とか聞かれるやつですね。
4 7
.....#G
.#.#.#.
.S#..#.
...#...
少しトリッキーですが、以下のようにして[][]intにして.を0、#を1、Sを2、Gを3にするところまでをやってくれます。文字列をint配列でとり、それを変換してます。
h, w := ni2()
mp := convidxi2s(nsi2s(h), map[string]int{".": 0, "#": 1, "S": 2, "G": 3})
何も用意していない状態で臨むとhでループ→文字列取得→wでループ→[]byteをstringに変換→判別というコードが必要になり手間となります。
出力
出力も必ずある上にパターンも似たり寄ったりなので準備しておくメリットが大きいです。
基本的な実装と全関数
以下が基本的な流れです。なぜバッファリングするかは冒頭で紹介した記事に述べられているので割愛しますが、こうしておくことで出力で困ったケースはほぼないはずです。なお、インタラクティブな問題でもflush()を出力タイミングで呼ぶことで対応できます。
var wtr = bufio.NewWriter(os.Stdout)
func main() {
defer flush()
o := 0
out(o)
}
func out(v ...interface{}) {
_, e := fmt.Fprintln(wtr, v...)
if e != nil {
panic(e)
}
}
func flush() {
e := wtr.Flush()
if e != nil {
panic(e)
}
}
入力よりは少ないですが、これだけ関数を作ってます。
自分用ライブラリのio関連部分
// vを出力
func out(v ...interface{})
// vをローカルのみ出力
func debug(v ...interface{})
// vをフォーマットして出力
func outf(f string, v ...interface{})
// vを改行なし(without ln)で出力
func outwoln(v ...interface{})
// intのスライスを半角スペース区切りで出力
func outis(sl []int)
// intのスライスを改行区切りで出力
func outisnr(sl []int)
// バッファにためた出力をflushする
func flush()
デバッグ
out(v)はfmt.Printlnで出せるものならなんでも出せるので当然デバッグにも使えます。ただ、消し忘れたり、消したはいいけどWAになってまたデバッグしたい時に困ったりすることがあります。以下のようなローカルでしか出力されないデバッグ関数を作っておき、残したまま提出してしまうのがおすすめです。
var debugFlg bool
func init() {
// ローカルの時のみ実行時に引数をつけてデバッグモードにする
if len(os.Args) > 1 && os.Args[1] == "i" {
debugFlg = true
}
}
func debug(v ...interface{}) {
if !debugFlg {
return
}
out(v...)
}
頻出パターン
スライスはスペース区切りで表示して出力するケースが多いです。連結させたり、別の文字が区切りの場合はさっとstrings.Joinの第二引数を書き換えます。
func outis(sl []int) {
r := make([]string, len(sl))
for i, v := range sl {
r[i] = itoa(v)
}
out(strings.Join(r, " "))
}
また改行しつつ表示もよくあります、
func outisnr(sl []int) {
for _, v := range sl {
out(v)
}
}
ちょうど二つを出したり、フォーマット調整するときはprintfを使うのですが、これも準備しておくと地味に便利です。fmt.Sprintfだと打つ量が倍になりますし、Sprintfだっけ？SPrintfだっけ？と忘れたりするので…。
func outf(f string, v ...interface{}) {
out(fmt.Sprintf(f, v...))
}
既存処理のラッパー
既存処理をエラー処理含めて包み込んでしまうのも地味に便利です。
例えば以下だとstrconv.Atoiと打って、エラーハンドリングをしなければいけないのがatoiと打つだけですっと終わります。
func atoi(s string) int {
i, e := strconv.Atoi(s)
if e != nil {
panic(e)
}
return i
}
func itoa(i int) string {
return strconv.Itoa(i)
}
引数をfloat64で渡すmath系の関数も厳密な数値が必要でない場合は役に立つことが多いです。例えば以下平方根を取る関数です。
func sqrt(i int) int {
return int(math.Sqrt(float64(i)))
}
以下のようにキャッシュしておくこともでき、意図せずTLEになることを防げたりもします。以下は2^nを求める関数です。
var pow2cache [64]int
func pow2(i int) int {
if pow2cache[i] == 0 {
pow2cache[i] = int(math.Pow(2, float64(i)))
}
return pow2cache[i]
}
よく使う関数があれば、それをラップして便利にできないか検討してもると良いでしょう。
型は極力intで扱う
既存処理のラッパーでも出てきましたが色々なものをintのまま扱えるようにしておくとストレスがないです。intとfloat64が混じるコードは非常に変換が手間で、コンパイルエラーがうわっと出てこれ全部にfloat64()とかint()とかせなあかんのか…となりがちです。
func sqrt(i int) int {
return int(math.Sqrt(float64(i)))
}
stringはstringのままでも扱えますが、使い方によって型がstring,[]byte,runeなど色々な型となるので煩雑です。intのスライスにしてしまい、最後に文字列が必要ならintのスライスから戻すとするとストレスがないです。
// stois("abcde",'a') -> []{0,1,2,3,4}
func stois(s string, baseRune rune) []int {
r := make([]int, len(s))
for i, v := range s {
r[i] = int(v - baseRune)
}
return r
}
// istos([]{0,1,2,3,4},'a') -> "abcde"
func istos(s []int, baseRune rune) string {
r := make([]byte, len(s))
for i, v := range s {
r[i] = byte(v) + byte(baseRune)
}
return string(r)
}
入力のところでも出ましたが、初手でintのスライスにしてしまうのも便利な局面が多いです。
as := nsis() // abcde -> []int{0,1,2,3,4}
スライス操作
スライスはよく扱うので、色々と便利な関数を準備できるかと思います。
intスライス作成
ただ作って初期値を入れるだけですが、使用頻度が非常に多いのでおすすめです。
func is(l int, def int) []int {
sl := make([]int, l)
for i := 0; i < l; i++ {
sl[i] = def
}
return sl
}
func i2s(l, m int, def int) [][]int {
sl := make([][]int, l)
for i := 0; i < l; i++ {
sl[i] = make([]int, m)
for j := 0; j < m; j++ {
sl[i][j] = def
}
}
return sl
}
ソート
よく出るのであると便利です。
func sorti(sl []int) {
sort.Sort(sort.IntSlice(sl))
}
func sortir(sl []int) {
sort.Sort(sort.Reverse(sort.IntSlice(sl)))
}
[2]intのスライスのソートはそれなりに出るのでこちらも便利です。
[2]intが[3]intや[4]intとかになっても引数の型をちょっと直すだけで対応できます。
type Sort2ArOptions struct {
keys []int
orders []sortOrder
}
type Sort2ArOption func(*Sort2ArOptions)
func opt2ar(key int, order sortOrder) Sort2ArOption {
return func(args *Sort2ArOptions) {
args.keys = append(args.keys, key)
args.orders = append(args.orders, order)
}
}
// sort2ar(sl,opt2ar(1,asc))
// sort2ar(sl,opt2ar(0,asc),opt2ar(1,asc))
func sort2ar(sl [][2]int, setters ...Sort2ArOption) {
args := &Sort2ArOptions{}
for _, setter := range setters {
setter(args)
}
sort.Slice(sl, func(i, j int) bool {
for idx, key := range args.keys {
if sl[i][key] == sl[j][key] {
continue
}
switch args.orders[idx] {
case asc:
return sl[i][key] < sl[j][key]
case desc:
return sl[i][key] > sl[j][key]
}
}
return true
})
}
mod演算
Goでやる限りストレスなくできないのではと思います。modintという型があって、演算子をオーバーライドして普通の演算かのように扱えれば楽なのですが…。
一応関数を作っていますが、関数を使わずに毎回%=modを書いていくこともあります。ただ、そうすると掛け算で巨大な数字通しを掛け合わせてしまいオーバーフローさせてしまうミスをよくやってしまいます。関数化したものを使うのも書くのが手間というデメリットがあるのですが有用だったりします。
func madd(a, b int) int {
a %= mod
b %= mod
a += b
if a >= mod || a <= -mod {
a %= mod
}
if a < 0 {
a += mod
}
return a
}
func mmul(a, b int) int {
a %= mod
b %= mod
return a * b % mod
}
func mdiv(a, b int) int {
a %= mod
if b <= 0 || b >= mod {
panic("invalid division")
}
return a * minvfermat(b, mod) % mod
}
二分探索、三分探索
二分探索と三分探索はそれぞれintとfloat64で使えるようにしています。二分探索はかなりの頻度で出題されます、三分探索は少ないものの一から書くのはそれなりに大変なので便利です。
https://github.com/gosagawa/atcoder/blob/bf0863a1d378c98eeff787665da1668e9ed6c140/_template/main.go#L724
あと、単純なintでの二分探索の場合は以下のようにlowerBoundとupperBoundを取る関数を用意しておくと二分探索の関数を意識する必要がなくなるので便利です。
// find x of sl[x] < v. return -1 if no lowerbound found
func lowerBound(v int, sl []int) int {
if len(sl) == 0 {
return -1
}
idx := bs(0, len(sl)-1, func(c int) bool {
return sl[c] < v
})
return idx
}
// find x of v < sl[x]. return len(sl) if no upperbound found
func upperBound(v int, sl []int) int {
if len(sl) == 0 {
return 0
}
idx := bs(0, len(sl)-1, func(c int) bool {
return sl[c] <= v
})
return idx + 1
}
sorted map / 平衡二分木
ここまでは時短を目的とした関数を主に紹介してきましたが、こちらは知っておかないと他の言語だと圧倒的に解かれているのに、Goだと全然解けないというトラップがあります。それがsorted mapです。Goのmapはソートされないのでソートされたmapを使いたい時に困ります。
AtCoder側もこの問題を把握していると思われ、ここだけ特例としてgithub.com/emirpasic/godsというライブラリが組み込まれて使えるようになっています。最低限github.com/emirpasic/gods/trees/redblacktreeに関しては使い方を覚えておくと良いでしょう。既にまとめてくださっている方がおり、詳しくは以下の記事が参考になるかと思います。
私はなんでここだけ外部ライブラリをよむのだと癪だったので、自力で赤黒木を作っています。ただ、アルゴリズム的には最速でなく使い方も分かりやすくはなっていないので上記の既存のライブラリを使う方が良いかと思います。
(appendix)時には自己流の関数ではなく、既存のよくできたライブラリを完コピすると良いかも
ここまで色々な自分専用関数を作ってきたのですが、大きめな関数になるほど手を入れたい時とか改良が大変だったり、更に改良した時にバグったりで混乱します。
遅延セグメント木は初めは自己流ライブラリを作っていたのですが色々な問題に対応できなかったりバグに苦しめられたりで、結局AtCoderライブラリからほぼ使用感が変わらないようにC++からGoに移植したものを作りました。
見ての通り、関数を呼び出すまでがゴツいのですが、世にチートシートがあったり、他の人のコードを参考にできたり、ライブラリ部分は作ってしまえば手を加えることを考える必要がなくなり非常に楽です。自分で一から作るのも勉強になりますが、すでに広く使われたライブラリの移植もおすすめです。C++からGoは言語のベースが近い部分もあり比較的移植しやすいように思います。
なお、AtCoderLibraryをそのまま全てGoに移植したものは現状なさそうです。有志の方が作られていたようで以下に一部存在してそうなのですが、アクティブではなさそうですね。今回欲しかった遅延セグメント木もありませんでした。
おそらくよほど今後Goでの競プロがない限り造られることはないでしょう…。
遅延セグメント木に関しては頑張って作って、同様のものは現状だと世になさそうなので使ってみていただければと思います。
https://github.com/gosagawa/atcoder/blob/bf0863a1d378c98eeff787665da1668e9ed6c140/_template/main.go#L1753
// 区間変更・区間最小値取得
op := func(a segstruct, b segstruct) segstruct {
return segstruct(min(int(a), int(b)))
}
e := func() segstruct {
return segstruct(inf)
}
id := func() segfstruct {
return segfstruct(inf)
}
mapping := func(f segfstruct, x segstruct) segstruct {
if f == id() {
return x
}
return segstruct(int(f))
}
compostion := func(f segfstruct, g segfstruct) segfstruct {
if i == id() {
return g
}
return f
}
lst := newlazysegtree(
n,
base,
op,
e,
mapping,
compostion,
id,
)
lst.applyrange(f, t, segfstruct(v))
iv := int(lst.get(i))
rv := int(lst.prod(l,r)
av := int(t.allprod()))
まとめ
以上、長くなりましたがGoで競技プログラムをやっていて得た関数周りのノウハウをまとめてみました。是非参考になれば幸いです！