https://leetcode.com/contest/leetcode-weekly-contest-48/problems/trim-a-binary-search-tree/

[1st version] not so good....

#include <iostream>
#include <queue>
using namespace std;

typedef struct TreeNode TreeNode;

struct TreeNode {
  int val;
  TreeNode *left;
  TreeNode *right;
  TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};
/**
 * Definition for a binary tree node.

 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */

class Solution {
public:
  TreeNode* trimBST(TreeNode* root, int L, int R) {
    if (!root) {
      return root;
    }
   
    TreeNode *lChild = trimBST(root->left,  L, R);
    TreeNode *rChild = trimBST(root->right, L, R);
    TreeNode *temp, *newRoot;
    newRoot = root;
    if (root->val < L || root->val > R) {
      newRoot = NULL;
      if (lChild) {
        temp    = root;
        newRoot = lChild;
        while(newRoot->right) {
          temp = newRoot;
          newRoot = newRoot->right;
        }
        if (newRoot != lChild) {
          temp->right    = NULL;
          newRoot->left  = lChild;
        }
        newRoot->right = rChild;
      } else if (rChild) {
        temp    = root;
        newRoot = rChild;
        while(newRoot->left) {
          temp = newRoot;
          newRoot = newRoot->left;
        }
        if (newRoot != rChild) {
          temp->left      = NULL;
          newRoot->right  = rChild;
        }
        newRoot->left = lChild;
      }
    } else {
      root->left  = lChild;
      root->right = rChild;
    }
    return newRoot;
  }
};

void inOrder(TreeNode* root)
{
  if (!root)
    return;

  if (root->left)
    inOrder(root->left);
  cout << root->val << "\n";
  if (root->right)
    inOrder(root->right);
}

void preOrder(TreeNode* root)
{
  if (!root)
    return;

  cout << root->val;
  cout << "\n";
  if (root->left)
    preOrder(root->left);
  if (root->right)
    preOrder(root->right);
}

void postOrder(TreeNode* root)
{
  if (!root)
    return;

  if (root->left)
    postOrder(root->left);
  if (root->right)
    postOrder(root->right);
  cout << root->val;
  cout << "\n";
}

int main()
{
  //int input[] = {2,2,5,-1,-1,5,7};
  int input[] = {3,2,4,1};
  int size = sizeof(input) / sizeof(int);
  cout << "size = " << size << "\n";

  TreeNode **Tree = new TreeNode*[size];

  for (int i = size - 1; i >= 0; i--) {
    if (input[i] != -1) {
      Tree[i] = new TreeNode(input[i]);
      int lChild = 2 * i + 1;
      int rChild = 2 * i + 2;
      Tree[i]->val = input[i];
      if (lChild < size && (input[lChild] != -1)) {
        Tree[i]->left  = Tree[lChild];
      }
      if (rChild < size && (input[rChild] != -1)) {
        Tree[i]->right  = Tree[rChild];
      }    
    }
  }
  cout << "preOrder\n";
  preOrder(Tree[0]);
  cout << "end\n\n";

  cout << "inOrder\n";
  inOrder(Tree[0]);
  cout << "end\n\n";

  cout << "postOrder\n";
  postOrder(Tree[0]);
  cout << "end\n\n";

  Solution solution;
  TreeNode *result = solution.trimBST(Tree[0], 2, 3);

  cout << "preOrder\n";
  preOrder(result);
  cout << "end\n\n";

  cout << "inOrder\n";
  inOrder(result);
  cout << "end\n\n";

  cout << "postOrder\n";
  postOrder(result);
  cout << "end\n\n";

  return 1;
}

[2nd version] simpler & shorter on

  TreeNode* trimBST2(TreeNode* root, int L, int R) {
    if (!root) {
      return root;
    }
    if (root->val < L)
      return trimBST2(root->right, L, R);
    else if (root->val > R)
      return trimBST2(root->left, L, R);
    else {
      root->right = trimBST2(root->right, L, R);
      root->left   = trimBST2(root->left, L, R);
      return root;
    }
  }

留言

這個網誌中的熱門文章

https://leetcode.com/contest/leetcode-weekly-contest-52/problems/longest-univalue-path/

https://leetcode.com/contest/leetcode-weekly-contest-52/problems/maximum-sum-of-3-non-overlapping-subarrays/