BSTIterator

Author: masx200

binary-search-tree-iterator/BSTIterator_test.go

@@ -0,0 +1,19 @@
+package binary_search_tree_iterator
+
+import (
+	"testing"
+
+	"github.com/masx200/leetcode-test/utils"
+	"gotest.tools/v3/assert"
+)
+
+var TreeNodeLeetCodeParse = utils.TreeNodeLeetCodeParse
+
+func TestBSTIterator(t *testing.T) {
+	iterator := Constructor(TreeNodeLeetCodeParse("[7, 3, 15, null, null, 9, 20]"))
+	res := []int{}
+	for iterator.HasNext() {
+		res = append(res, iterator.Next())
+	}
+	assert.DeepEqual(t, res, []int{3, 7, 9, 15, 20})
+}

binary-search-tree-iterator/index.go

@@ -0,0 +1,74 @@
+package binary_search_tree_iterator
+
+import (
+	"context"
+	serialize_and_deserialize_binary_tree "github.com/masx200/leetcode-test/serialize-and-deserialize-binary-tree"
+	"runtime"
+)
+
+type TreeNode = serialize_and_deserialize_binary_tree.TreeNode
+
+type BSTIterator struct {
+	generator chan int
+	value     int
+	done      bool
+}
+
+func Constructor(root *TreeNode) BSTIterator {
+	ctx, cancel := context.WithCancel(context.Background())
+	gen := InOrderIterator(ctx, root)
+
+	value, done := <-gen
+	bstiterator := BSTIterator{generator: gen, value: value, done: done}
+	runtime.SetFinalizer(&bstiterator, func(_ *BSTIterator) { cancel() })
+	return bstiterator
+}
+
+func (b *BSTIterator) Next() int {
+	val := b.value
+	b.value, b.done = <-b.generator
+	return val
+}
+
+func (b *BSTIterator) HasNext() bool {
+	return b.done
+}
+
+func InOrderIterator(ctx context.Context, root *TreeNode) (gen chan int) {
+	gen = make(chan int)
+
+	go func() {
+
+		defer close(gen)
+		if root == nil {
+
+			return
+		}
+		YieldAll(ctx, gen, InOrderIterator(ctx, root.Left))
+		YieldOne(
+			ctx, gen, root.Val)
+
+		YieldAll(ctx, gen, InOrderIterator(ctx, root.Right))
+
+	}()
+
+	return
+}
+func YieldAll(ctx context.Context, target chan int, source chan int) {
+
+	for v := range source {
+		YieldOne(
+			ctx, target, v)
+
+	}
+}
+func YieldOne(ctx context.Context, target chan int, source int) {
+
+	select {
+	case <-ctx.Done():
+		return
+	default:
+		target <- source
+	}
+
+}