プログラミング
記事内に商品プロモーションを含む場合があります

Go言語の ゴルーチンとチャネルでfor~rangeによる繰り返しを実現する方法

tadanori

Go言語でPythonのイテレーターやジェネレータに似た関数をチャネル(channel)とゴルーチン(goroutine)を使って実装する方法を解説します。チャネルを使うことで、for e := range f() のようにして要素を生成する関数を使ったループ処理が記述できます

はじめに

最近、競プログラム(AtCoder)の問題を解くときに組み合わせ列挙のプログラムでTLEしてしまい組み合わせ列挙の自作関数を見直しました。

この関数はゴルーチンとチャネルを利用したジェネレータとして実装していたのですが、最近ごルーチンを全然作っていなかったので、久しぶりのジェネレータ関数の作り直しに苦戦しました。

この記事では、自身の復習を兼ねて、for~rangeでチャネルをジェネレーターとして使う方法に焦点をあてて解説したいと思います。

ジェネレータの実装

ジェネレータ/イテレータとは

Pythonではイテレータ(iterator)とジェネレータ(generator)があり、ジェネレータはイテレータの一種のような扱いです。

  • イテレータ:反復して要素を取り出すことのできるインタフェース
  • ジェネレータ:イテレータの一種で、yield文を使ったものを指すことが多い

Go言語でチャネルを使った実装は、ジェネレータに近いと思いますので、ここではジェネレータと呼ぶことにします。

正式名称が分かりませんでした。たぶん、Go言語でもジェネレータで良いと思います。

ジェネレータは、例えば以下のようにfor~rangeで利用できるものを指します。

for i := range generator {
}

ジェネレータでは要素を作りながら処理ができるので、「あらかじめ要素を全て用意しておいてループさせる」のと比較して以下の点でメリットがあります

  • ループ数が処理により変化する/無限にループさせたい
  • 要素数が多く、あらかじめ用意しておくとメモリを消費してしまう

このような場合は、ジェネレータを使うのが便利です。

チャネルによるfor~range処理(基本形)

Go言語のfor~range文は、スライスやマップの要素を繰り返し処理するための文であり、rangeキーワードを使用して、スライスやマップの要素を繰り返し処理できます。

for~range文でチャネルを使ってゴルーチンで要素を生成する方法が、ジェネレータの基本的な作り方になります。

以下は、ジェネレータの基本形です

The Go Playgroundを使って動作確認を行うことができます。試してみてください。

package main

import "fmt"

func main() {
    // チャネルを生成
    ch := make(chan int)

    // 処理
    go func() {
        for i := 0; i < 10; i++ {
            ch <- i
        }
        close(ch)
    }()

    // チャネル for~range
    for i := range ch {
        fmt.Println(i)
    }
}
実行結果
0
1
2
3
4
5
6
7
8
9

処理の流れは以下になります

  1. チャネルを作成
  2. チャネルを使って0〜9の値を返すゴルーチンを起動
  3. for i := range chで、チャネル経由で要素を受け取る

ゴルーチンで要素を生成してチャネルに入れ、それをfor文で取り出します。

rangeの後の部分を関数化する

基本形を少し変形して、ジェネレータを関数化してみます。

ジェネレータを関数化した例は、以下になります。

関数化するメリットはgen(10)のように引数で要素数を渡すことです。

package main

import "fmt"

func main() {
	// ジェネレータ関数
	gen := func(n int) chan int {
		ch := make(chan int)
		go func() {
			for i := 0; i < n; i++ {
				ch <- i
			}
			close(ch)
		}()
		return ch
	}

	// ジェネレータの使用
	for i := range gen(10) {
		fmt.Println(i)
	}
}

先ほどと違うのは、gen関数自体の戻り値がchan intとチャネルになっていることです。チャネルを返す関数なので、for文でチャネルとして利用することが可能となります。

gen関数の中では、ゴルーチンを起動し、そこで要素の生成とチャネルへの入力を行っています。

また、gen関数をmainの外に置く場合は、以下のようになります。関数の宣言の書き方が少し違うだけで、先ほどの例とほぼ同じです。

package main

import "fmt"


func gen(num int) (c chan int) {
  c = make(chan int)

  go func() {
    for i := 0; i < num; i++ {
	   c <- i
	}
	close(c)
  }()
  
  return
}

func main() {
    // チャネル for~range
    for i := range gen(10) {
        fmt.Println(i)
    }
}

実用的な例(組み合わせ列挙)

0〜10のループでは、いまいちジェネレータの利点にピンとこないかと思います。

ここでは、実用的なジェネレータの例として、組み合わせを列挙するジェネレータを作ってみたいと思います。

組み合わせ列挙

組み合わせ列挙は、nCkといった記号で表記される組み合わせを全て列挙するものです。

例えば、{1,2,3,4,5}という整数の集合から2つの異なる整数の組み合わせを列挙する場合を考ええます。

この場合、対象の集合は{1, 2, 3, 4, 5}であり、この中から2つを組み合わせたものを列挙することになります。

この条件を満たす組み合わせは、{(1, 2), (1, 3), (1, 4), (1, 5), (2, 3), (2, 4), (2, 5), (3, 4), (3, 5), (4, 5)}の10通りとなります(組み合わせでは(1,2)と(2,1)は同じなので片方だけ列挙します)

これを列挙するのが組み合わせ列挙です

組み合わせをグラフにすると以下のような感じになります。「1を選んだら、次は2以降から選ぶ」、「2を選んだら3以降から選ぶ」といった具合です。

再帰を使った実装

順列列挙の動きは再帰関数を使えば比較的簡単に実装できます。

以下は実装例です

次の説明のために、変わった実装をしています。例えばpatはあらかじめ2個確保していますが、本来はappendしていくのが普通です

アルゴリズム的には、先ほど説明したものを素直に実装しています。

package main

import "fmt"

func rec(pos, start, n, k int, pat []int, f func([]int)) {
	// k個選んでいれば、それを出力する
	if pos == k {
		f(pat)
		return
	}
	// 選んでいない場合は、追加して再帰
	// 次に選べるのは、startからn-1までの値のいずれか
	for i := start; i < n; i++ {
		pat[pos] = i                  // posに選んだ数字をセットして
		rec(pos+1, i+1, n, k, pat, f) // pos, startを1つずつ進める
	}
}

func main() {
	f := func(x []int) {
		fmt.Println(x)
	}
	pat := make([]int, 2)
	rec(0, 0, 5, 2, pat, f)
}

結果は、以下のようになり10個の列挙が行われていることがわかります

実行結果
[0 1]
[0 2]
[0 3]
[0 4]
[1 2]
[1 3]
[1 4]
[2 3]
[2 4]
[3 4]

このプログラムでは、関数fに生成された組み合わせを渡すように実装されています(このような作り方をコールバックと呼びます)

このように実装すれば、組み合わせのパターンに合わせて好きな処理を行わせることができます。

でも、少し面倒です。

ジェネレータを使うと、for文を使って書くことができるようになり、すっきりします。

ゴルーチンとチャネルによりfor~rangeで使えるようにする

ゴルーチンを使って書き直したバージョンがこちらです。

package main

import "fmt"

func comb(n, k int) (c chan []int) {
	pat := make([]int, k)
	c = make(chan []int, 1)

	var rec func(pos, start int)

	rec = func(pos, start int) {
		// k個選んでいれば、それを出力する
		if pos == k {
			tmp := make([]int, k)
			copy(tmp, pat)
			c <- tmp
			return
		}
		// 選んでいない場合は、追加して再帰
		// 次に選べるのは、startからn-1までの値のいずれか
		for i := start; i < n; i++ {
			pat[pos] = i    // posに選んだ数字をセットして
			rec(pos+1, i+1) // pos, startを1つずつ進める
		}
	}
	go func() {
		defer close(c)
		rec(0, 0)
	}()

	return
}

func main() {
	for e := range comb(5, 2) {
		fmt.Println(e)
	}
}

関数を利用するときに以下のように簡潔に使うことができるので便利です。

for文で呼び出せるのはメリットです。プログラムがすっきりします

	for e := range comb(5, 2) {
		fmt.Println(e)
	}

以下、プログラムを見ていきます。

次の部分は引数が少し省略されている以外は、再帰関数のバージョンと同じです。ただ、チャネルに生成した組み合わせを送る部分で、生成したパターンをコピーしてそれを渡しています。

コピーしないと、次の生成で上書きされて結果がおかしくなります。ゴルーチンが並行動作していることに注意してください。

	rec = func(pos, start int) {
		// k個選んでいれば、それを出力する
		if pos == k {
			tmp := make([]int, k)
			copy(tmp, pat)
			c <- tmp
			return
		}
		// 選んでいない場合は、追加して再帰
		// 次に選べるのは、startからn-1までの値のいずれか
		for i := start; i < n; i++ {
			pat[pos] = i    // posに選んだ数字をセットして
			rec(pos+1, i+1) // pos, startを1つずつ進める
		}
	}

これをゴルーチンで呼び出すために、以下のようなラッパー関数を用意し、ゴルーチンとして呼び出しています。

	go func() {
		defer close(c)
		rec(0, 0)
	}()

こうすることで、再帰関数部分をほとんど変えなくて良いので楽です。

まとめ

以上、ゴルーチンとチャネルを使ってGo言語でジェネレータを作成する方法を解説しました。

簡単に説明していますが、久しぶりにゴルーチンを扱うと、チャネルを先にクローズしてしまったり、デットロック警告がでたりと結構手間がかかりました。

この記事が参考になれば幸いです。

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

記事URLをコピーしました