main.cpp

// Source: https://leetcode.com/problems/trim-a-binary-search-tree
// Title: Trim a Binary Search Tree
// Difficulty: Medium
// Author: Mu Yang <http://muyang.pro>

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Given the `root` of a binary search tree and the lowest and highest boundaries as `low` and `high`, trim the tree so that all its elements lies in `[low, high]`. Trimming the tree should **not** change the relative structure of the elements that will remain in the tree (i.e., any node's descendant should remain a descendant). It can be proven that there is a **unique answer**.
//
// Return the root of the trimmed binary search tree. Note that the root may change depending on the given bounds.
//
// **Example 1:**
// https://assets.leetcode.com/uploads/2020/09/09/trim1.jpg
//
// ```
// Input: root = [1,0,2], low = 1, high = 2
// Output: [1,null,2]
// ```
//
// **Example 2:**
// https://assets.leetcode.com/uploads/2020/09/09/trim2.jpg
//
// ```
// Input: root = [3,0,4,null,2,null,null,1], low = 1, high = 3
// Output: [3,2,null,1]
// ```
//
// **Constraints:**
//
// - The number of nodes in the tree is in the range `[1, 10^4]`.
// - `0 <= Node.val <= 10^4`
// - The value of each node in the tree is **unique**.
// - `root` is guaranteed to be a valid binary search tree.
// - `0 <= low <= high <= 10^4`
//
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include <algorithm>

using namespace std;

struct TreeNode {
  int val;
  TreeNode* left;
  TreeNode* right;
  TreeNode() : val(0), left(nullptr), right(nullptr) {}
  TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
  TreeNode(int x, TreeNode* left, TreeNode* right) : val(x), left(left), right(right) {}
};

// DFS
class Solution {
 public:
  TreeNode* trimBST(TreeNode* root, int low, int high) {
    if (!root) return nullptr;

    // Only right child
    if (root->val < low) {
      return trimBST(root->right, low, high);
    }

    // Only left child
    if (root->val > high) {
      return trimBST(root->left, low, high);
    }

    // Keep both child
    root->left = trimBST(root->left, low, high);
    root->right = trimBST(root->right, low, high);
    return root;
  }
};