2023.12.15

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

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

tadanori

AtCoderの解説などをみていると、C++のlower_bound, upper_boundを使った解法をよく目にします。この記事では、Go言語でlower_bound, upper_boundライクな関数の実装の仕方を解説します。

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

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

Contents

C++のlower_bound, upper_bound関数
Go言語のSortパッケージ
lowerBound, upperBoundを実装する
- 少し高度な使い方
まとめ

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の問題を解いていると欲しくなったので、この関数をGo言語用に作ってみました。

lowerBound, upperBoundを組み込んだ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関数はテンプレートとして用意しておいた方が便利です