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

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

Aru

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

はじめに

AtCoderの競技プログラミング問題を解く際、組み合わせを列挙するプログラムでTLE(時間超過)してしまい、自作の組み合わせ列挙関数を見直すことになりました。

この、組み合わせ列挙の関数はゴルーチンとチャネルを利用して実装していたのですが、久しぶりにゴルーチンを使ったジェネレータ関数を作り直すことになりする、予想以上に苦労しました。

この記事では、私自身のゴルーチンを使ったジェネレータの復習も兼ねて、Go言語におけるチャネルを使った関数の実装方法、そしてそれをfor~rangeでどのように活用できるかに焦点をあてて解説していきたいと思います。

ジェネレータの実装

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

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

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

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

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

ちなみにGo言語も1.23よりイテレータとyieldが実装されたようですので、そのうちそちらを使ったバージョンにもチャレンジしたいと思います。

ここでのジェネレータは、例えば以下のように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言語でジェネレータを作成する方法を解説しました。

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

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

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

ABOUT ME
ある/Aru
ある/Aru
IT&機械学習エンジニア/ファイナンシャルプランナー(CFP®)
専門分野は並列処理・画像処理・機械学習・ディープラーニング。プログラミング言語はC, C++, Go, Pythonを中心として色々利用。現在は、Kaggle, 競プロなどをしながら悠々自適に活動中
記事URLをコピーしました