Go言語でAtCoder｜スタック、キュー、優先キューの実装方法

Aru

本記事では、Go言語で、Stack、Queue、Priority Queue（優先度付きキュー）を実装する方法について解説します。この3つのデータ構造は、AtCoderなどの競技プログラムでも実務でも高頻度で利用するものです。

Contents

Go言語でAtCoder
スタック（Stack）・キュー（Queue）の作り方
優先キュー（Priority Queue）
まとめ

Go言語でAtCoder

私は、Go言語でAtCoderのコンテストに参加しています（最近はリアルタイムでの参加ではなく、過去問に取り組むことが多いです）。

いつも感じるのは、AtCoderの問題を解くときに必要な情報がGo言語には少ないということです。Go言語で参加する人が少ないため、情報が不足している理由だと思います。それならということでGo言語でAtCoderに参加する人向けに情報を発信しようというのがこの記事の趣旨になります。

一見難しそうに思える部分も、慣れてしまえばシンプルです。

その他のAtCoderに役立つ記事の一覧

AtCoderで役立つ記事一覧（Python, Go言語）

スタック（Stack）・キュー（Queue）の作り方

Go言語でスタックとキューを実装するのは、めちゃくちゃ簡単です。

一応、標準ライブラリにcontainer/list パッケージがありますが、そんなものを使う必要はなくスライスだけで実現できます。

Stack

後に入れたデータから先に取り出すデータ構造がスタックです。スタックはLast-In-First-Out（LIFO）と呼ばれることもあります。

スライスでは、appendでデータの末尾にデータを追加できるので、削除する場合にlen(x)-1を削除するようにすればスタックとして利用できます。

Go言語のスライスは、結構高速なので、スタックと利用することができます。

// スライスを作成
x := make([]int, 0)
// 1を末尾に追加
x = append(x, 1)
fmt.Println(x)
// 2を末尾に追加
x = append(x, 2)
fmt.Println(x)
// 末尾の２を削除
x = x[:len(x)-1]
fmt.Println(x)

[1]
[1 2]
[1]

まとめると、スタックの追加、削除手順は以下の通りです。

追加
x = append(x, n)としてスライスの末尾にnを追加
削除
x = x[:len(x)-1]としてスライスの末尾を削除

Queue

先に入れたデータから先に取り出すデータ構造がキューです。キューはFirst-In-Fisrt-Out（FIFO）と呼ばれることもあります。

Go言語でキューを作る場合は、スライスから削除するデータを先頭にして、次のデータから末尾までを元の変数に代入します。具体的には、下記のようにx = x[1:]と書きます。

スタックとほぼ同様のコードです。こちらも十分高速に動作します。

// スライスを作成
x := make([]int, 0)
// 1を末尾に追加
x = append(x, 1)
fmt.Println(x)
// 2を末尾に追加
x = append(x, 2)
fmt.Println(x)
// 先頭を削除
x = x[1:]
fmt.Println(x)

[1]
[1 2]
[2]

まとめると、スタックの追加、削除手順は以下の通りです。

追加
x = append(x, n)としてスライスの末尾にnを追加
削除
x = x[1:]としてスライスの先頭を削除

構造体をキューに入れる

キューやスタックに、構造体を入れたい場合もあります。この場合は以下のように記述できます。

以下は、２変数（a,b）からなる構造体をキューで出し入れする例です。書き方は、先ほどとほとんど変わりません。

type data struct {
	a, b int
}

x := make([]data, 0)
x = append(x, data{1, 2})
fmt.Println(x)
x = append(x, data{3, 4})
fmt.Println(x)
x = x[1:]
fmt.Println(x)

[{1 2}]
[{1 2} {3 4}]
[{3 4}]

優先キュー（Priority Queue）

Go言語で優先キューを使う場合は、container/heapパッケージを使います。importにcontainer/heapを追加してください。なお、VSCodeなどでGo言語のプラグインを使っている場合は、自動で追加されるので特に気にする必要はありません。

ジェネリスクに対応した優先キューのGo言語用のコードは以下にあります。こちらも参考にしてください。

ヒープと優先キューの作り方をGolangで学ぼう｜ジェネリクス対応版

import (
	"container/heap"
)

まず、優先キューで使うデータ型を定義します。汎用性があるように、pqi構造体を定義しました。構造体はxだけしかメンバがないので、実質整数を格納しているのと同じです。キューに入れるデータを増やしたい場合は、ここにyなど変数と型を追加していってください。

また、優先キューのソート順は、Less()で定義します。例では、xの昇順になります。ここを書き換えることで様々な条件でソートすることが可能です。そういう意味では、他の言語より使いやすいかもしれません。

例えば、xについて昇順、xが同値についてはyについて降順などの優先キューも定義することができます。

Len, Swap, Less, Push, Popの関数が定義されていることが、heapを使うための条件になります。このあたりの書き方はGoっぽいです。

type pqi struct{ x int }

type priorityQueue []pqi

func (pq priorityQueue) Len() int            { return len(pq) }
func (pq priorityQueue) Swap(i, j int)       { pq[i], pq[j] = pq[j], pq[i] }
func (pq priorityQueue) Less(i, j int) bool  { return pq[i].x < pq[j].x }
func (pq *priorityQueue) Push(x interface{}) { *pq = append(*pq, x.(pqi)) }
func (pq *priorityQueue) Pop() interface{} {
	x := (*pq)[len(*pq)-1]
	*pq = (*pq)[0 : len(*pq)-1]
	return x
}

上記はテンプレートにしておくと良いと思います。

上記を書いておけば、使い方は以下の通りです。

pqをpriorityQueue型の空の変数として定義します。要素の追加は、heap.Push()で、削除は、heap.Pop()で行います。heap.Pop()では、先頭の要素を取り出すことができますが、下に書いたような、x.(pqi)という書き方が必要になります。

これはGo言語の型アサーションというものですが、毎回書くのも面倒です。

先頭は、pq[0]でアクセスできますので、自分はpq[0]でアクセスした後、heap.Pop()で削除しています。この方が、覚えることが少ないのでおすすめです。

	pq := priorityQueue{}

	heap.Push(&pq, pqi{10})
	heap.Push(&pq, pqi{20})
	heap.Push(&pq, pqi{1})

	fmt.Println(pq[0], len(pq))

	q := heap.Pop(&pq).(pqi)
	fmt.Println(q.x)

	fmt.Println(pq[0], len(pq))