Go言語による各種データ構造の扱い方(やや競プロer向け)

こんにちは株式会社スマートショッピング ソフトウェアエンジニアの葛西伸樹です。
今年の1月ごろから競技プログラミングのコンテスト(AtCoder)にGo言語で参加し始めたのですが、使用人口も少なく情報が見つからずなかなか苦戦しましたので現在わかっている範囲でGo言語による各種データ構造の扱い方についてまとめました。
指摘やコメントがあればぜひお願いします！

この記事で説明すること/しないこと

この記事では競技プログラミングでよく使われるデータ構造のGo言語による扱い方とおおよその実行速度について説明していきます。
今回説明するデータ構造は以下の通りです。

• 配列
• 動的配列
• キュー/スタック
• 優先度付きキュー
• HashMap/HashSet
• TreeMap/TreeSet

また

• 各種データ構造自体の説明
• 各種データ構造を使用した具体的なアルゴリズム

については説明を省略します。

早速各データ構造の説明を始めたいと思います。

配列

Go言語でも他の言語と同様に配列が用意されていますがsizeが固定であり、次に説明する動的配列であるsliceとそこまで速度の差がないためあまり使う機会はないと思います。

// 要素数を指定して宣言をする
var a [1e7]int
for i := 0; i < 1e7; i++ {
	// データのPush
	a[i] = i
	// データのGet
	_ = a[i]
}
// AtCoderでの実行速度
// 実行時間 34ms
// メモリ 80712KB

動的配列

Go言語では要素数を動的に変更できる配列としてsliceというものが用意されています。
宣言の方法としては

1. sliceのリテラルを使う方法
2. make関数を使う方法

があります。

// 1. sliceのリテラルを使う方法
a := []int{1, 2, 3, 4, 5}
// 2. make関数を使う方法
b := make([]int, 5, 10)

make関数の第一、第二引数ではそれぞれsliceのsizeとcapacityを指定することができます。
細かい説明は省略しますが、sliceに要素を追加してsizeがcapacityを超える時にメモリの再確保が発生して処理速度に影響が出てしまうので宣言時には取りうる最大のsizeをcapacityとして指定することが望ましいです。

// capacityを指定しないで初期化
a := make([]int, 0)
for i := 0; i < 1e7; i++ {
	// データのPush
	a = append(a, i)
	// データのGet
	_ = a[i]
}
// AtCoderでの実行速度
// 実行時間 269ms
// メモリ 362920KB

// capacityを指定して初期化
b := make([]int, 0, 1e7)
for i := 0; i < 1e7; i++ {
	// データのPush
	b = append(b, i)
	// データのGet
	_ = b[i]
}
// AtCoderでの実行速度
// 実行時間 38ms
// メモリ 80712KB

キュー/スタック

Go言語では標準ライブラリでキューやスタック専用のデータ型が用意されていないので先ほど説明したsliceが使われることが多いです。

queue := make([]int, 0, 1e7)
for i := 0; i < 1e7; i++ {
	// データのPush
	queue = append(queue, i)
}
for i := 0; i < 1e7; i++ {
	// 先頭からデータをPop
	_ = queue[0]
	queue = queue[1:]
}
// AtCoderでの実行速度
// 実行時間 38ms
// メモリ 80708KB

stack := make([]int, 0, 1e7)
for i := 0; i < 1e7; i++ {
	// データのPush
	stack = append(stack, i)
}
for i := 0; i < 1e7; i++ {
	// 末尾からデータをPop
	_ = stack[len(stack)-1]
	stack = stack[:len(stack)-1]
}
// AtCoderでの実行速度
// 実行時間 44ms
// メモリ 80708KB

ただしsliceは先頭に要素を追加する場合メモリの再確保が発生して処理速度が遅くなってしまいます。

stack := make([]int, 0, 1e7)
for i := 0; i < 1e7; i++ {
	// データを先頭にPush
	stack = append([]int{i}, stack...)
}
for i := 0; i < 1e7; i++ {
	// 先頭からデータをPop
	_ = stack[0]
	stack = stack[1:]
}
// AtCoderでの実行速度
// 実行時間 TLE(>10514ms)
// メモリ -

なのでsliceをDequeのように前後から要素を追加したり削除するようなデータ構造として利用することは難しいので、標準ライブラリのcontainer/listを使う方がいいでしょう。

deque := list.New()
for i := 0; i < 1e7; i++ {
	// データを先頭にPush
	deque.PushFront(i)
}
for i := 0; i < 1e7; i++ {
	// 先頭からデータをPop
	_ = deque.Front().Value.(int)
	deque.Remove(deque.Front())
}
// AtCoderでの実行速度
// 実行時間 1801ms
// メモリ 576320KB

優先度付きキュー

Go言語では標準ライブラリのcontainer/heapで優先度付きキューが用意されています。
container/heapでは指定されたインターフェースを満たす型を実装することで用意されたメソッドを使うことができます。

// IntHeapの実装
type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// 実際の処理
func main() {
	ih := &IntHeap{}
	heap.Init(ih)
	for i := 0; i < 1e7; i++ {
		// データのPush
		// 性能を見るために、1e4で割った余りをPushしている
		heap.Push(ih, i%1e4)
	}
	for i := 0; i < 1e7; i++ {
		// データのPop
		_ = heap.Pop(ih)
	}
}
// AtCoderでの実行速度
// 実行時間 3993ms
// メモリ 405380KB

また下記はダイクストラなどで使うようなEdge(グラフの辺)の優先度付きキューの実装例です。

type Edge struct {
	To     int // 移動先のNode
	Weight int // コスト
}
type EdgeHeap []Edge

func (h EdgeHeap) Len() int           { return len(h) }
func (h EdgeHeap) Less(i, j int) bool { return h[i].Weight < h[j].Weight }
func (h EdgeHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *EdgeHeap) Push(x interface{}) {
	*h = append(*h, x.(Edge))
}

func (h *EdgeHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func main() {
	eh := &EdgeHeap{}
	heap.Init(eh)
	for i := 0; i < 1e7; i++ {
		// データのPush
		// 性能を見るために、1e4で割った余りをPushしている
		heap.Push(eh, Edge{To: i, Weight: i % 1e4})
	}
	for i := 0; i < 1e7; i++ {
		// データのPop
		_ = heap.Pop(eh).(Edge)
	}
}
// AtCoderでの実行速度
// 実行時間 4695ms
// メモリ 658724KB

HashMap/HashSet

Go言語ではmapというHashMapが用意されています。
仕組みについてはhi6okuniさんの記事が参考になるのでぜひ読んでみてください！
宣言の方法としては

1. mapのリテラルを使う方法
2. make関数を使う方法

があります。

// 1. mapのリテラルを使う方法
a := map[int]int{
	1: 1,
	2: 2,
	3: 3,
}
// 2. make関数を使う方法
b := make(map[int]int, 3)

mapの場合もsliceと同様にmake関数の引数でcapacityが指定できるのでこちらも取りうる最大のsizeを指定することが望ましいです。

// capacityを指定しないで初期化
a := make(map[int]int)
for i := 0; i < 1e7; i++ {
	// データのPush
	a[i] = i
}
for i := 0; i < 1e7; i++ {
	// データのGet
	_ = a[i]
}
// AtCoderでの実行速度
// 実行時間 2413ms
// メモリ 638728KB

// capacityを指定して初期化
b := make(map[int]int, 1e7)
for i := 0; i < 1e7; i++ {
	// データのPush
	b[i] = i
}
for i := 0; i < 1e7; i++ {
	// データのGet
	_ = b[i]
}
// AtCoderでの実行速度
// 実行時間 1564ms
// メモリ 322572KB

またGo言語にはHashSet専用の型は用意されていないのでmapで代用することが多いです。
mapをHashSetとして扱う方法として

1. mapのvalueを空の構造体(struct{})にする
2. mapのvalueをboolにする

の2つがあります。
それぞれ1.は使用するメモリを最小にできる、2.は存在しない場合falseが返ってくるのでそのまま判定に使用することでコードの記述量を減らせるというメリットがあります。
競技プログラミングにおいてメモリがネックになることがあまりないのでコードを書く速度を上げるために私は2.を選択することが多いです。

// valueを空のstructとして宣言
setA := make(map[int]struct{})
// データのPush
setA[1] = struct{}{}
// データの存在確認
if _, ok := setA[1]; ok {
	// 存在する場合の処理
}

// valueをboolとして宣言
setB := make(map[int]bool)
// データのPush
setB[1] = true
// データの存在確認
if setB[1] {
	// 存在する場合の処理
}

TreeMap/TreeSet

Go言語では標準でTreeMap/TreeSetを提供していないため外部ライブラリであるredblacktreeを使用する必要があります。
AtCoderではこちらのライブラリを使用することができるため自分で実装する必要はありません。

rbt := redblacktree.NewWithIntComparator()
for i := 0; i < 1e7; i++ {
	// データのPut
	rbt.Put(i, i)
}
for i := 0; i < 1e7; i++ {
	// データのGet
	_, _ = rbt.Get(i)
}
// AtCoderでの実行速度
// 実行時間 6862ms
// メモリ 902436KB

まとめ

競技プログラミングでよく使われるデータ構造のGo言語による扱い方をまとめました。Go言語でコンテストに参加している方やこれから使ってみたいという方の参考になれば幸いです。
また弊社では競プロ部としてみんなで問題を解いたりアルゴリズムに関する知見の共有をして切磋琢磨していますのでその活動に関してもいつか記事にしようと思っています！