Neunomizuの日記

俺だけが俺だけじゃない

# LeetCode Medium 173. Binary Search Tree Iterator

tags: leetcode

問題

イデア

二分木のイテレータを作れという問題です.

next()hasNext()というメソッドが時間,空間計算量を平均でそれぞれ$O(1)$,$O(h)$となるようにするという制約があります

解法1

木を配列に変換して,そこからノードを取るという風にします

配列にはinorder,つまり小さい順にノードを入れます.

計算量

ノードの数を$N$として,木の高さを$h$をします

  • 時間計算量
    • 木を配列に変換するのに$O(N)$
    • next()hasNext()に$O(1)$`
  • 空間計算量
    • 再帰の深さから$O(h)$
    • next()hasNext()に$O(1)$
    • ノードを格納する配列に$O(N)$

C++

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class BSTIterator {
  vector<int> sortedNodes;    
  int index;
  
public:
  BSTIterator(TreeNode* root) {
    index = -1;
    _inorder(root);
  }

  /** @return the next smallest number */
  int next() {
    return sortedNodes.at(++index);
  }

  /** @return whether we have a next smallest number */
  bool hasNext() {
    return index + 1 < sortedNodes.size();
  }
  
private:
  void _inorder(TreeNode* root) {
    if (!root)
      return;
    
    _inorder(root->left);
    sortedNodes.push_back(root->val);
    _inorder(root->right);
  }
};

/**
 * Your BSTIterator object will be instantiated and called as such:
 * BSTIterator* obj = new BSTIterator(root);
 * int param_1 = obj->next();
 * bool param_2 = obj->hasNext();
 */

解法2

上の解法では二分木のノードの数だけ空間が必要なのでこれをより効率的にするためにスタックを使います.

最初にスタックにinorderでノードを入れます.つまり,値の小さいノードからスタックに積んでいきます.

  • next()
    • スタックの先頭のノードを最小値として返す
    • その際に先頭のノードはポップする
    • そして,先頭の右の子にの左部分木をスタックに積む
  • hasNext
    • 現在のノードの数が0より大きければまだノードが取り出せるので`true

計算量

  • 時間計算量
    • 最初にノードをスタックに積むのに$O(h)$
    • next()hasNext()は共に$O(1)$
  • 空間計算量
    • スタックに積むノードは木の高さに依存するので$O(h)$
    • next()は$O(h)$が必要

C++

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class BSTIterator {
  stack<TreeNode*> nodeStack;
public:
  BSTIterator(TreeNode* root) {
    _leftMostInorder(root);
  }

  /** @return the next smallest number */
  int next() {
    auto topNode = nodeStack.top();
    nodeStack.pop();
    if (topNode->right)
      _leftMostInorder(topNode->right);
    return topNode->val;
  }

  /** @return whether we have a next smallest number */
  bool hasNext() {
    return nodeStack.size() > 0;
    }
private:
  void _leftMostInorder(TreeNode* root) {
    while (root) {
      nodeStack.push(root);
      root = root->left;
    }
  }
};

/**
 * Your BSTIterator object will be instantiated and called as such:
 * BSTIterator* obj = new BSTIterator(root);
 * int param_1 = obj->next();
 * bool param_2 = obj->hasNext();
 */