aboutsummaryrefslogtreecommitdiff
path: root/add_row_tree.cpp
blob: 54ebfb4da1d8c988a8233877c6a007151423561d (plain)
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
#include "leetcode.h"

class Solution {
public:
  TreeNode *addOneRow(TreeNode *root, int val, int depth) {
    if (depth-- == 1)
      return new TreeNode(val, root, nullptr);
    queue<TreeNode *> q;
    q.push(root);
    while (!q.empty()) {
      bool addLevel = depth-- == 1;
      size_t sz = q.size();
      for (size_t _ = 0; _ < sz; _++) {
        TreeNode *curr = q.front();
        q.pop();
        if (addLevel) {
          curr->left = new TreeNode(val, curr->left, nullptr);
          curr->right = new TreeNode(val, nullptr, curr->right);
          continue;
        }
        if (curr->left)
          q.push(curr->left);
        if (curr->right)
          q.push(curr->right);
      }
    }
    return root;
  }
};