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

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

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

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

C++のlower_bound, upper_bound関数

lower_bound と upper_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言語の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言語）

ある/Aru

IT＆機械学習エンジニア/ファイナンシャルプランナー（CFP®)

XContact

専門分野は並列処理・画像処理・機械学習・ディープラーニング。プログラミング言語はC, C++, Go, Pythonを中心として色々利用。現在は、Kaggle, 競プロなどをしながら悠々自適に活動中