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

Go言語でC++のlower_bound, upper_boundを実装する

Aru

C++のlower_bound, upper_boundは、ソートされたコンテナを二分探索により高速に検索する関数です。AtCoderなどの競技プログラミングでよく利用されるこれらの関数ですが、Go言語(golang)には標準ライブラリとしてはsort.Search(二分探索)しか用意されていません。この記事では、C++のlower_bound, upper_boundの機能を持つ関数をSearchを使って実装します。また、ジェネリクス対応の二分探索の実装方法について紹介します。

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

AtCoderで役立つ記事一覧(Python, Go言語)
AtCoderで役立つ記事一覧(Python, Go言語)

C++のlower_bound, upper_bound関数

lower_boundupper_bound は、C++のSTL(標準テンプレートライブラリ)に含まれる関数で、主にソートされた配列やコンテナ内で特定の値の位置を探すために利用します。

lower_bound

例えば、以下のような配列でlower_boundで4を探すと、最初の4の位置(2の次の4の位置)を返します。

{1, 2, 4, 4, 5, 6, 7}

upper_bound

一方、upper_boundで以下の配列から4を探すと、4より大きい最初の要素の5の位置を返します。

{1, 2, 4, 4, 5, 6, 7}

C++ではSTLとして用意されている

二分探索をするだけなので、自分で書いても良いですが、C++ではライブラリとして用意されているため、関数を呼び出すだけで済みます。一方、Go言語には、これらの関数は用意されていません。

AtCoderの問題を解いているうちに、「毎回作る」のが面倒になったので、lowerBound, upperBoundをGo言語向けに作成してみました。

lowerBound, upperBoundを組み込んだAtCoder用のテンプレートも公開しています。詳しくは、以下の記事を参照してください。

Go言語でAtCoder参加時のテンプレートと注意すべきポイント
Go言語でAtCoder参加時のテンプレートと注意すべきポイント
AtCoderの問題について

AtCoder Beginner Contestの問題を解いていると、C++の標準関数を前提として難易度設定されている印象を受けます。lower_bound, upper_boundがライブラリにあれば一行で簡単に実装できるから難易度低めみたいな感じです。
Go言語で参加する場合は、この手の関数はテンプレート化して用意しておくことをおすすめします。

Go言語のSortパッケージ

実はGo言語では、二分探索を行うための関数がsortパッケージに用意されています

このsort.Search関数を利用すると、lowerBound, upperBoundを簡単に実装することができます。

ここでは、sort.Search関数について簡単に説明します。

sort.Search関数のインタフェースは以下の通りです。

func Search(n int, f func(i int) bool) int

インタフェースが少し複雑で、[0, n)の範囲を探索し、関数fがtrueとなる最初の位置を返します。

例えば、以下のようなコードの場合、最初に2よりも大きくなる3が返ります。

	dat := []int{1, 2, 3, 4, 5, 6, 7}
	x := 2
	pos := sort.Search(len(dat), func(i int) bool {
		return dat[i] > x
	})
	fmt.Println(pos, dat)
出力結果
2 [1 2 3 4 5 6 7]

最初にtrueとなる位置なので、datを降順に並べえると、以下のような検索も可能です。

	dat := []int{7, 6, 5, 4, 3, 2, 1}
	x := 2
	pos := sort.Search(len(dat), func(i int) bool {
		return dat[i] < x
	})
	fmt.Println(pos, dat)
出力結果
6 [7 6 5 4 3 2 1]

このように、Search関数を使えば二分探索を簡単に記述することができます。

実は、Go言語ではlowerBound, upperBoundより汎用的なライブラリが用意されているわけです。これを使って、lowerBound, upperBoundを実装します。

lower_boundではなく、lowerBoundにしているのは、Go言語の命名規則に合わせるためです。

lowerBound, upperBoundを実装する

以下のコードがlowerBound, upperBoundのコードになります。

違いは、return a[i] >= xなのかreturn a[i] > x なのかです。Search関数を使って毎回書いてもよいですが、関数を使った方が楽ですので、活用してください。

func lowerBound(a []int, x int) int {
	idx := sort.Search(len(a), func(i int) bool {
		return a[i] >= x
	})
	return idx
}

func upperBound(a []int, x int) int {
	idx := sort.Search(len(a), func(i int) bool {
		return a[i] > x
	})
	return idx
}

使い方は以下のようになります。lowerBoundでは、3番目(0インデックス)の4を、upperBoundでは5番目の5の位置を返していることが確認できます。

	dat := []int{1, 2, 4, 4, 5, 6, 7}
	fmt.Println(lowerBound(dat, 4), upperBound(dat, 4))
出力結果
2, 4

C++のlower_bound, upper_boundではイテレータを返しますが、今回実装した関数では配列のインデックスを返します注意してください。

少し高度な使い方

sort.Searchを使えば、構造体のlowerBoundなどをとることも可能です。

例えば、数値と文字列からなる構造体を定義して、lowerBoundをとることも可能です。

type pair struct {
	idx int
	s   string
}

func lowerBound(a []pair, x int) int {
	idx := sort.Search(len(a), func(i int) bool {
		return a[i].idx >= x
	})
	return idx
}

func main() {
	dat := []pair{
		pair{1, "aaa"},
		pair{2, "xyz"},
		pair{3, "xcc"},
		pair{4, "min"},
	}

	pos := lowerBound(dat, 4)
	fmt.Println(pos, dat[pos])
}

sort.Searchの使い方も覚えておくと便利です。

とはいえ、lowerBound, upperBound関数はテンプレートとして用意しておいた方が便利です

ジェネリクス対応版の二分探索

先に書いたように、sort.Searchを直接使えば構造体の二分探索を行うことが可能ですが、ここでは、ジェネリクス対応版の二分探索(binSearch)の関数を作成してみました。

コードは以下になります。基本は、sort.Searchを使うのと同じです。

func binSearch[T comparable](a []T, x T, less func(T, T) bool) int {
	l, r := 0, len(a)
	for l+1 != r {
		m := (l + r) / 2
		if less(a[m], x) {
			r = m
		} else {
			l = m
		}
	}
	return r
}

上記の関数では、sort.Searchメソッドを使わずに二分探索を直接実装しています。二分探索はそこまで難しくないので、記述方法は覚えておいて損はないです。

lower_boundとして利用する

整数型のlower_boundとして利用する場合は、以下のようにします(x >= y)。

	dat := []int{1, 2, 4, 4, 5, 6, 7}
	fmt.Println(binSearch[int](dat, 4, func(x, y int) bool {
		return x >= y
	}))

upper_boundとして利用する

整数型のupper_boundとして利用する場合は、以下のようにします(x > y)。

	dat := []int{1, 2, 4, 4, 5, 6, 7}
	fmt.Println(binSearch[int](dat, 4, func(x, y int) bool {
		return x > y
	}))

まとめ

C++のlower_bound, upper_boundに相当する関数を、Go言語で実装する方法について解説しました。

他にも、AtCoderで活用できる内容を記事にしています。

あわせて読みたい
AtCoderで役立つ記事一覧(Python, Go言語)
AtCoderで役立つ記事一覧(Python, Go言語)

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

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