c++のstd::set::lower_boundとupper_boundをGoで作ってみました

螺旋本の Part3 第16章 16.13「線分交差問題」（対応するオンラインジャッジは AOJ CGL_6_A）を Go で解こうとしたところ、
螺旋本に載っている C++ のサンプルコードで std::set::lower_bound / upper_bound を使用しており、
Go には同等の関数が無かったため、自分で実装しました。

注意

• 二分探索木で作っています
• イテレータではなく、Nodeを返します
• 四角い車輪の再発明です
  ※ この lower_bound / upper_bound は 二分探索木 (BST) 用の実装です。
  []int やスライスに対して使える関数ではありません。
  （sort.Search() を使えばスライスに対してバイナリサーチできます）

Go版lower_boundのコード

// x 以上の最小の key をもつノードへのポインタを返す。
// 見つからなければ nil を返す。
func lowerBoundNode(n *Node, x int) *Node {
	var res *Node
	for n != nil {
		if n.key >= x {
			// 条件を満たしたので「候補」として覚えておく
			res = n
			// もっと小さい key で条件を満たすノードが左側に
			// あるかもしれないので、左部分木を探索
			n = n.left
		} else {
			// n.key < x のときは、左側は全部 x 未満なので捨てて
			// 右部分木へ
			n = n.right
		}
	}
	return res
}

func (t *Tree) lowerBound(x int) *Node {
	return lowerBoundNode(t.root, x)
}

t.lowerBound(x) で、key >= x となる最小の *Node を返します。

Go版upper_boundのコード

// x より大きい最小の key をもつノードへのポインタを返す。
// 見つからなければ nil を返す。
func upperBoundNode(n *Node, x int) *Node {
	var res *Node
	for n != nil {
		if n.key > x {
			// 条件を満たしたので「候補」として覚えておく
			res = n
			// もっと小さい key で条件を満たすノードが左側に
			// あるかもしれないので、左部分木を探索
			n = n.left
		} else {
			// n.key <= x のときは、左側は全部 x 以下なので捨てて
			// 右部分木へ
			n = n.right
		}
	}
	return res
}

func (t *Tree) upperBound(x int) *Node {
	return upperBoundNode(t.root, x)
}

t.upperBound(x)で、key > x となる最初の*Nodeを返します。
lowerBoundと大体同じです。

コード全部

package main

import (
	"fmt"
)

type Node struct {
	key    int
	parent *Node
	left   *Node
	right  *Node
}

type Tree struct {
	root *Node
}

func (t *Tree) insert(k int) {
	x := t.root
	var y *Node
	for x != nil {
		y = x
		if k < x.key {
			x = x.left
		} else if k > x.key {
			x = x.right
		} else {
			return // 重複は追加しない
		}
	}
	newNode := &Node{key: k, parent: y}
	if y == nil {
		t.root = newNode
	} else if k < y.key {
		y.left = newNode
	} else {
		y.right = newNode
	}
}

// lower_bound(key >= x の最小ノード)
func lowerBoundNode(n *Node, x int) *Node {
	var res *Node
	for n != nil {
		if n.key >= x {
			res = n
			n = n.left
		} else {
			n = n.right
		}
	}
	return res
}

func (t *Tree) lowerBound(x int) *Node {
	return lowerBoundNode(t.root, x)
}

// upper_bound(key > x の最小ノード)
func upperBoundNode(n *Node, x int) *Node {
	var res *Node
	for n != nil {
		if n.key > x {
			res = n
			n = n.left
		} else {
			n = n.right
		}
	}
	return res
}

func (t *Tree) upperBound(x int) *Node {
	return upperBoundNode(t.root, x)
}

// 中間順巡回で木の中身を昇順表示する
func inorder(n *Node) {
	if n == nil {
		return
	}
	inorder(n.left)
	fmt.Printf("%d ", n.key)
	inorder(n.right)
}

func main() {
	// テストデータ（BST に入れる値）
	values := []int{10, 5, 15, 3, 7, 13, 18}

	t := &Tree{}
	for _, v := range values {
		t.insert(v)
	}

	fmt.Print("Tree inorder = ")
	inorder(t.root)
	fmt.Println("\n")

	tests := []int{0, 4, 5, 6, 7, 10, 14, 15, 19}

	for _, x := range tests {
		fmt.Printf("x = %2d | ", x)

		lb := t.lowerBound(x)
		if lb != nil {
			fmt.Printf("lowerBound = %2d, ", lb.key)
		} else {
			fmt.Printf("lowerBound = nil, ")
		}

		ub := t.upperBound(x)
		if ub != nil {
			fmt.Printf("upperBound = %2d\n", ub.key)
		} else {
			fmt.Printf("upperBound = nil\n")
		}
	}
}

参考

お気楽 Go 言語プログラミング入門(二分探索木)
プログラミングコンテスト攻略のためのアルゴリズムとデータ構造(Amazon)
https://cpprefjp.github.io/reference/set/set/lower_bound.html
https://cpprefjp.github.io/reference/set/set/upper_bound.html