LCR 055-二叉搜索树迭代器

Raphael Liu Lv10

实现一个二叉搜索树迭代器类BSTIterator ,表示一个按中序遍历二叉搜索树(BST)的迭代器:

  • BSTIterator(TreeNode root) 初始化 BSTIterator 类的一个对象。BST 的根节点 root 会作为构造函数的一部分给出。指针应初始化为一个不存在于 BST 中的数字,且该数字小于 BST 中的任何元素。
  • boolean hasNext() 如果向指针右侧遍历存在数字,则返回 true ;否则返回 false
  • int next()将指针向右移动,然后返回指针处的数字。

注意,指针初始化为一个不存在于 BST 中的数字,所以对 next() 的首次调用将返回 BST 中的最小元素。

可以假设 next() 调用总是有效的,也就是说,当调用 next() 时,BST 的中序遍历中至少存在一个下一个数字。

示例:

**输入**
inputs = ["BSTIterator", "next", "next", "hasNext", "next", "hasNext", "next", "hasNext", "next", "hasNext"]
inputs = [[[7, 3, 15, null, null, 9, 20]], [], [], [], [], [], [], [], [], []]
**输出**
[null, 3, 7, true, 9, true, 15, true, 20, false]

**解释**
BSTIterator bSTIterator = new BSTIterator([7, 3, 15, null, null, 9, 20]);
bSTIterator.next();    // 返回 3
bSTIterator.next();    // 返回 7
bSTIterator.hasNext(); // 返回 True
bSTIterator.next();    // 返回 9
bSTIterator.hasNext(); // 返回 True
bSTIterator.next();    // 返回 15
bSTIterator.hasNext(); // 返回 True
bSTIterator.next();    // 返回 20
bSTIterator.hasNext(); // 返回 False

提示:

  • 树中节点的数目在范围 [1, 105]
  • 0 <= Node.val <= 106
  • 最多调用 105hasNextnext 操作

进阶:

  • 你可以设计一个满足下述条件的解决方案吗?next()hasNext() 操作均摊时间复杂度为 O(1) ,并使用 O(h) 内存。其中 h 是树的高度。

注意:本题与主站 173 题相同: <https://leetcode-cn.com/problems/binary-search-tree-
iterator/>

前言

根据二叉搜索树的性质,不难发现,原问题等价于对二叉搜索树进行中序遍历。因此,我们可以使用中序遍历的思路来解决这一问题。

方法一:扁平化

我们可以直接对二叉搜索树做一次完全的递归遍历,获取中序遍历的全部结果并保存在数组中。随后,我们利用得到的数组本身来实现迭代器。

[sol1-C++]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class BSTIterator {
private:
void inorder(TreeNode* root, vector<int>& res) {
if (!root) {
return;
}
inorder(root->left, res);
res.push_back(root->val);
inorder(root->right, res);
}
vector<int> inorderTraversal(TreeNode* root) {
vector<int> res;
inorder(root, res);
return res;
}

vector<int> arr;
int idx;
public:
BSTIterator(TreeNode* root): idx(0), arr(inorderTraversal(root)) {}

int next() {
return arr[idx++];
}

bool hasNext() {
return (idx < arr.size());
}
};
[sol1-Java]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class BSTIterator {
private int idx;
private List<Integer> arr;

public BSTIterator(TreeNode root) {
idx = 0;
arr = new ArrayList<Integer>();
inorderTraversal(root, arr);
}

public int next() {
return arr.get(idx++);
}

public boolean hasNext() {
return idx < arr.size();
}

private void inorderTraversal(TreeNode root, List<Integer> arr) {
if (root == null) {
return;
}
inorderTraversal(root.left, arr);
arr.add(root.val);
inorderTraversal(root.right, arr);
}
}
[sol1-Golang]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
type BSTIterator struct {
arr []int
}

func Constructor(root *TreeNode) (it BSTIterator) {
it.inorder(root)
return
}

func (it *BSTIterator) inorder(node *TreeNode) {
if node == nil {
return
}
it.inorder(node.Left)
it.arr = append(it.arr, node.Val)
it.inorder(node.Right)
}

func (it *BSTIterator) Next() int {
val := it.arr[0]
it.arr = it.arr[1:]
return val
}

func (it *BSTIterator) HasNext() bool {
return len(it.arr) > 0
}
[sol1-JavaScript]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
var BSTIterator = function(root) {
this.idx = 0;
this.arr = [];
this.inorderTraversal(root, this.arr);
};

BSTIterator.prototype.next = function() {
return this.arr[this.idx++];
};

BSTIterator.prototype.hasNext = function() {
return this.idx < this.arr.length;
};

BSTIterator.prototype.inorderTraversal = function(root, arr) {
if (!root) {
return;
}
this.inorderTraversal(root.left, arr);
arr.push(root.val);
this.inorderTraversal(root.right, arr);
};
[sol1-C]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
typedef struct {
int* res;
int size;
int idx;
} BSTIterator;

int getTreeSize(struct TreeNode* root) {
if (root == NULL) {
return 0;
}
return 1 + getTreeSize(root->left) + getTreeSize(root->right);
}

void inorder(int* ret, int* retSize, struct TreeNode* root) {
if (root == NULL) {
return;
}
inorder(ret, retSize, root->left);
ret[(*retSize)++] = root->val;
inorder(ret, retSize, root->right);
}

int* inorderTraversal(int* retSize, struct TreeNode* root) {
*retSize = 0;
int* ret = malloc(sizeof(int) * getTreeSize(root));
inorder(ret, retSize, root);
return ret;
}

BSTIterator* bSTIteratorCreate(struct TreeNode* root) {
BSTIterator* ret = malloc(sizeof(BSTIterator));
ret->res = inorderTraversal(&(ret->size), root);
ret->idx = 0;
return ret;
}

int bSTIteratorNext(BSTIterator* obj) {
return obj->res[(obj->idx)++];
}

bool bSTIteratorHasNext(BSTIterator* obj) {
return (obj->idx < obj->size);
}

void bSTIteratorFree(BSTIterator* obj) {
free(obj->res);
free(obj);
}

复杂度分析

  • 时间复杂度:初始化需要 O(n) 的时间,其中 n 为树中节点的数量。随后每次调用只需要 O(1) 的时间。

  • 空间复杂度:O(n),因为需要保存中序遍历的全部结果。

方法二:迭代

除了递归的方法外,我们还可以利用栈这一数据结构,通过迭代的方式对二叉树做中序遍历。此时,我们无需预先计算出中序遍历的全部结果,只需要实时维护当前栈的情况即可。

[sol2-C++]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class BSTIterator {
private:
TreeNode* cur;
stack<TreeNode*> stk;
public:
BSTIterator(TreeNode* root): cur(root) {}

int next() {
while (cur != nullptr) {
stk.push(cur);
cur = cur->left;
}
cur = stk.top();
stk.pop();
int ret = cur->val;
cur = cur->right;
return ret;
}

bool hasNext() {
return cur != nullptr || !stk.empty();
}
};
[sol2-Java]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class BSTIterator {
private TreeNode cur;
private Deque<TreeNode> stack;

public BSTIterator(TreeNode root) {
cur = root;
stack = new LinkedList<TreeNode>();
}

public int next() {
while (cur != null) {
stack.push(cur);
cur = cur.left;
}
cur = stack.pop();
int ret = cur.val;
cur = cur.right;
return ret;
}

public boolean hasNext() {
return cur != null || !stack.isEmpty();
}
}
[sol2-Golang]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
type BSTIterator struct {
stack []*TreeNode
cur *TreeNode
}

func Constructor(root *TreeNode) BSTIterator {
return BSTIterator{cur: root}
}

func (it *BSTIterator) Next() int {
for node := it.cur; node != nil; node = node.Left {
it.stack = append(it.stack, node)
}
it.cur, it.stack = it.stack[len(it.stack)-1], it.stack[:len(it.stack)-1]
val := it.cur.Val
it.cur = it.cur.Right
return val
}

func (it *BSTIterator) HasNext() bool {
return it.cur != nil || len(it.stack) > 0
}
[sol2-JavaScript]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
var BSTIterator = function(root) {
this.cur = root;
this.stack = [];
};

BSTIterator.prototype.next = function() {
while (this.cur) {
this.stack.push(this.cur);
this.cur = this.cur.left;
}
this.cur = this.stack.pop();
const ret = this.cur.val;
this.cur = this.cur.right;
return ret;
};

BSTIterator.prototype.hasNext = function() {
return this.cur !== null || this.stack.length;
};
[sol2-C]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
typedef struct {
struct TreeNode* cur;
struct StackTreeNode* stk[128];
int stkSize;
} BSTIterator;

BSTIterator* bSTIteratorCreate(struct TreeNode* root) {
BSTIterator* ret = malloc(sizeof(BSTIterator));
ret->cur = root;
ret->stkSize = 0;
return ret;
}

int bSTIteratorNext(BSTIterator* obj) {
while (obj->cur != NULL) {
obj->stk[(obj->stkSize)++] = obj->cur;
obj->cur = obj->cur->left;
}
obj->cur = obj->stk[--(obj->stkSize)];
int ret = obj->cur->val;
obj->cur = obj->cur->right;
return ret;
}

bool bSTIteratorHasNext(BSTIterator* obj) {
return obj->cur != NULL || obj->stkSize;
}

void bSTIteratorFree(BSTIterator* obj) {
free(obj);
}

复杂度分析

  • 时间复杂度:显然,初始化和调用 hasNext() 都只需要 O(1) 的时间。每次调用 next() 函数最坏情况下需要 O(n) 的时间;但考虑到 n 次调用 next() 函数总共会遍历全部的 n 个节点,因此总的时间复杂度为 O(n),因此单次调用平均下来的均摊复杂度为 O(1)。

  • 空间复杂度:O(n),其中 n 是二叉树的节点数量。空间复杂度取决于栈深度,而栈深度在二叉树为一条链的情况下会达到 O(n) 的级别。

 Comments
On this page
LCR 055-二叉搜索树迭代器