- #include <iostream>
- #include <queue>
- #include <unordered_map>
- #include <unordered_set>
- #include <vector>
- using namespace std;
- class TreeNode {
- TreeNode *left;
- TreeNode *right;
- int val;
- public:
- TreeNode(TreeNode *l, TreeNode *r, int v) : left(l), right(r), val(v) {}
- TreeNode(int v) : left(nullptr), right(nullptr), val(v) {}
- TreeNode *get_left() { return left; }
- TreeNode *get_right() { return right; }
- int get_val() { return val; }
- };
- // replace with a single val
- TreeNode *copyReplaceValHelper(TreeNode *node, int v) {
- if (!node) return nullptr;
- TreeNode *left_copied = copyReplaceValHelper(node->get_left(), v);
- TreeNode *right_copied = copyReplaceValHelper(node->get_right(), v);
- return new TreeNode(left_copied, right_copied, v);
- }
- TreeNode *copyReplaceVal(TreeNode *tree, int v) {
- return copyReplaceValHelper(tree, v);
- }
- // replace with inorder
- TreeNode *copyReplaceInorderHelper(TreeNode *node, const vector<int> &vals, int &idx) {
- if (!node) return nullptr;
- TreeNode *left_copied = copyReplaceInorderHelper(node->get_left(), vals, idx);
- int val_idx = idx++;
- TreeNode *right_copied = copyReplaceInorderHelper(node->get_right(), vals, idx);
- return new TreeNode(left_copied, right_copied, vals[val_idx]);
- }
- TreeNode *copyReplaceInorder(TreeNode *tree, const vector<int> &vals) {
- int idx = 0;
- return copyReplaceInorderHelper(tree, vals, idx);
- }
- // replace with level order
- TreeNode *copyReplaceLevelOrderHelper(TreeNode *node, unordered_map<TreeNode *, int> &hmap, const vector<int> &vals) {
- if (!node) return nullptr;
- TreeNode *left_copied = copyReplaceLevelOrderHelper(node->get_left(), hmap, vals);
- TreeNode *right_copied = copyReplaceLevelOrderHelper(node->get_right(), hmap, vals);
- return new TreeNode(left_copied, right_copied, vals[hmap[node]]);
- }
- TreeNode *copyReplaceLevelOrder(TreeNode *tree, const vector<int> &vals) {
- if (!tree) return nullptr;
- // create node->idx mapping
- unordered_map<TreeNode *, int> hmap;
- queue<TreeNode *> q({tree});
- int idx = 0;
- while (!q.empty()) {
- TreeNode *node = q.front();
- q.pop();
- hmap[node] = idx++;
- if (node->get_left()) q.push(node->get_left());
- if (node->get_right()) q.push(node->get_right());
- }
- // copy
- return copyReplaceLevelOrderHelper(tree, hmap,vals);
- }
- void printTree(const std::string &prefix, TreeNode *node, bool isLeft) {
- if (node != nullptr) {
- std::cout << prefix;
- std::cout << (isLeft ? "├──" : "└──");
- // print the value of the node
- std::cout << node->get_val() << std::endl;
- // enter the next tree level - left and right branch
- printTree(prefix + (isLeft ? "│ " : " "), node->get_left(), true);
- printTree(prefix + (isLeft ? "│ " : " "), node->get_right(), false);
- }
- }
- void printTree(TreeNode *node) {
- printTree("", node, false);
- }
- int main() {
- // size = 6
- TreeNode *input = new TreeNode(new TreeNode(new TreeNode(1), new TreeNode(3), 2), new TreeNode(new TreeNode(5), nullptr, 6), 4);
- printTree(input);
- // replace val
- cout << "replace val " << endl;
- printTree(copyReplaceVal(input, 0));
- // replace inorder
- cout << "replace inorder " << endl;
- printTree(copyReplaceInorder(input, {1, 3, 5, 7, 9, 2}));
- // replace level order
- cout << "replace level order " << endl;
- printTree(copyReplaceLevelOrder(input, {1, 2, 3, 4, 5, 6}));
- }
复制代码 |