Go言語でAtCoder｜Stack, Queue, Priority Queueの書き方

tadanori

Contents

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

Go言語でAtCoder

Go言語でAtCoderのコンテストに参加しています（最近は、リアルタイム参加ではなく、後追いですが）。

Go言語でAtCoderに参加している人が少ないので、情報が少なめです。

ということで、私の方でも情報を書いていこうと思います。

慣れてしまうとなんてことないものもありますが、これからGo言語を使って参加しようと考えている方の参考になれば幸いです。

その他のAtCoderに役立つ記事の一覧は以下にあります。

スタック（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]

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]

構造体をキューに入れる

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

以下は、２変数（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言語のプラグインを使っている場合は、自動で追加されるので特に気にする必要はありません。

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))