🎨

AtCoderをGoで解くためにあると便利な関数

2020/09/27に公開

はじめに

具体的に作りたいものはまだないけどなんとなく触ってみたい言語があるっていう人には競プロは良い題材だと思います．その点においてAtCoderは対応言語が非常に多く，週1単位でコンテストがあり，過去問にもチャレンジすることができ，他参加者の提出が閲覧でき，そしてなにより楽しいと非常に便利なサービスです．
かくいう私もなんとなくGo言語を触りたいと思ったときにAtCoderのコンテストにGoを使って参加することで少しはGoを書けるようになったと自負（錯覚）しております．

この記事はそんな筆者がAtCoderの問題をGoで解くにあたって「この関数は必須だよね」，「この関数はあると便利だよね」といったヘルパー関数を必要度ごとに区切って紹介する記事です．
また，筆者はGo以前ではPythonを使用していたため，ヘルパー関数の必要度がPython基準となっています．(筆者がGoを使用し始めた理由の1つにPythonでは想定解でもTLEになってしまう問題に当たることが増えてきたというのもあります)

今回紹介する関数が含まれているテンプレート(普段私が使用しているもの)はこのリポジトリにまとめてあります．適宜参照してください．

必須レベル

一覧(必須レベル)

標準入力系
型変換系(string -> int変換など)
abs,max,min,mod

標準入力系

何も考えずに標準入力を使用するならfmtパッケージのScanfが使用できるが

var a int
fmt.Scanf("%d",&a)

この例のように2行必要となるうえに記述量も多い．なのでbufioのReaderを使用します．(他の提出者を見るとScannerを使用している方が多い印象です)

const BUFSIZE = 10000000
var rdr *bufio.Reader

func main() {
	// BUFSIZEは大きなint値を指定しておく
	rdr = bufio.NewReaderSize(os.Stdin, BUFSIZE)
	solve()
}

readline(戻り値 string)とreadIntSlice(戻り値 []int)の定義．readlineにてbufのcapを16としているがGoのSliceはappend操作でcapが足りなくなった場合に2倍されていくので速度的には問題ありません.^[1]

// 一行をstringで読み込み
func readline() string {
	buf := make([]byte, 0, 16)
	for {
		l, p, e := rdr.ReadLine()
		if e != nil {
			fmt.Println(e.Error())
			panic(e)
		}
		buf = append(buf, l...)
		if !p {
			break
		}
	}
	return string(buf)
}

// 一行を[]intで読み込み
func readIntSlice() []int {
	slice := make([]int, 0)
	lines := strings.Split(readline(), " ")
	for _, v := range lines {
		// s2iはstringをintに変換する関数(後述)
		slice = append(slice, s2i(v))
	}
	return slice
}

あとこれは完全に個人の好み(Pythonに頭を侵食されてしまった筆者の好み)なのですがGo言語でも以下のように標準入力から得た値をそれぞれの変数に格納したい．
10 20という行をint型の数値として読み込む場合:

// このように書きたい
a,b := readInt()

// このようには書きたくない
row := readIntSlice()
a,b := row[0],row[1]

そのため不格好ですが以下のような関数を用意します.

// s2iはstringをintに変換する関数(後述)
func readint() int {
	return s2i(readline())
}

func readint2() (int, int) {
	lines := strings.Split(readline(), " ")
	return s2i(lines[0]), s2i(lines[1])
}

func readint3() (int, int, int) {
	lines := strings.Split(readline(), " ")
	return s2i(lines[0]), s2i(lines[1]), s2i(lines[2])
}
...

これを使用すると以下のように書ける.

a,b := readint2()

ただこれはかなり不格好に見えてしまうんですよね...Pythonだったら

a,b = [int(v) for v in input().split()]

みたいに書けるんですけどね．(Pythonと比較するのは自分でもどうかと思いますが)
他にスマートな書き方があれば教えてください．

型変換系

string -> int(int -> string)だったりbool -> int(int -> bool)などの変換．コンテスト中にいちいちstrconv.Atoiなどの戻り値を確認してエラー処理をするのが面倒なので私には必須.

// string <-> int
func s2i(s string) int {
	v, ok := strconv.Atoi(s)
	if ok != nil {
		panic("Faild : " + s + " can't convert to int")
	}
	return v
}

func i2s(i int) string {
	return strconv.Itoa(i)
}

// bool <-> int
func b2i(b bool) int {
	if b {
		return 1
	}
	return 0
}

func i2b(i int) bool {
  return i != 0
}

abs,max,min,mod

このあたりを用意した理由は筆者の頭がPython脳になっているためです．競プロでは手放せません．それぞれ個別に説明します.

abs

説明が不要かもしれませんがint型の絶対値を返す関数です．float型の絶対値が欲しくなったら引数と戻り値の型を変更すれば使えます．

func abs(v int) int {
	if v < 0 {
		return -v
	}
	return v
}

max,min

max,minは引数の中から最大，最小の値を返す処理です．
いちいちfor文まわすのも面倒ですし管理するべき変数が多くなるので関数にしています．

func min(values ...int) int {
	ret := values[0]
	for _, v := range values {
		if ret > v {
			ret = v
		}
	}
	return ret
}

func max(values ...int) int {
	ret := values[0]
	for _, v := range values {
		if ret < v {
			ret = v
		}
	}
	return ret
}

mod

個人的には一番大事.この処理に何度か煮え湯を飲まされています．
Go言語でのmod(%)は負の値を返すことがあります．たしかC++のmod処理も同様だったと思います．
一方で一部言語(Pythonなど)では負の値が帰ることはありません．

-3 % 5
// Goの場合 : -3
// Pythonの場合 : 2

競プロで行うmod処理は正の値が帰ってきてくれたほうが嬉しいことが多いのでPythonと同じような挙動を行うmod関数を導入しています．

func mod(x, y int) int {
	m := x % y
	if m < 0 {
		return m + y
	}
	return m
}
mod(-3,5) // 2が帰ってくる

DPで完璧に解けたと思った問題がWAにされたことが2回ほどあるので皆様気をつけてください．

あると便利

一覧(あると便利)

pow
lcm,gcm

pow

べき乗計算を行う．mathパッケージのPowは引数，戻り値の方がfloat64なのでintのべき乗がほしいっていうときにキャストしなきゃいけなくてちょっと不便ですし，記述量も多くなるので以下の関数を使用しています．

func pow(x, y int) int {
	return int(math.Pow(float64(x), float64(y)))
}
// pow(2,3) // 8

lcm,gcm

最小公倍数，最大公約数を算出する．

func gcd(v1, v2 int) int {
	if v1 > v2 {
		v1, v2 = v2, v1
	}
	for v1 != 0 {
		v1, v2 = v2%v1, v1
	}
	return v2
}

func lcm(v1, v2 int) int {
	return v1 * v2 / gcd(v1, v2)
}

ただ他の関数と違ってアルゴリズムの中身になるので気にする場合は入れないほうがいいかもしれないですね．

関数以外の話題

GoのSliceについて

GoのSliceを他言語の双方向リストのような扱い方をするのはおすすめできません．
具体的にはPopやInsertなどの処理が以下のような書き方をしなければいけません．

[]intのSliceに処理を行う場合 :

// Pop 末尾から要素(x)を取り出す 
x := slice[n-1]
slice = slice[0 : n-1]

// Insert xを先頭に挿入
slice = append([]int{x},...slice)

Pop操作やInsert操作はQueueを用いた幅優先探索やStackを用いた深さ優先探索を行う場合に頻出しますが，Sliceでは上記のように扱いづらいです．(うろ覚えですが上記のInsert操作は非常に遅かったと記憶しています．)

なのでこれらの操作にはGoの双方向リスト実装であるcontainer/listを使用することをおすすめします．(参照)[https://golang.org/pkg/container/list/]
これを使用するとlistの操作は以下のようにかけます．

list(l)に対して処理を行う場合 :

// Pop 末尾から要素(x)を取り出す 
x := l.Remove(l.Back()).(int)

// Insert xを先頭に挿入
x := l.PushFront(x)

l.Removeの戻り値の型がinterface{}なのでキャストする必要がある点に注意が必要です．
もちろん双方向リストなので先頭から要素の取り出し(l.Remove(l.Front()))や末尾に要素を追加(l.PushBack(x))も行えます．

int64について

今回紹介した関数ではint型を使用しています．
競プロでは32bitには収まらないのでint64を明示的に使用したほうが良いのではないかと思われるかもしれませんが，私個人の意見としてはintで十分だと思います．
というのもGoのint型はシステムが32bitコンピュータであれば32bit,64bitコンピュータであれば64bitとして扱われます．参照
最近のPCはほとんどが64bitでしょうし，今までint型を使用してAtCoderのジャッジサーバに採点されてきましたがintが32bitとして扱われてWAになることもなかったのでおそらく採点サーバも64bitコンピュータなんだと思います．
型名を入力するときに64と打つのも冗長な気がするので今のところはintで良いと思います．

脚注

今のところはここでTLEになったことはないです．本当なら計測するべきなので後々計測を行いたいですね． ↩︎