leet_code/103.zigzagLevelOrder.cc at master · pagict/leet_code · GitHub

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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
#include <vector>

#include "leetcode_common_struct.h"

/// 3 ms Beats 56.97 % of users with C++
/// 12.24 MB Beats 98.58 % of users with C++

inline TreeNode *pop_next(std::vector<TreeNode *> &nodes, int level_end,
                          int is_stack) {
  if (nodes.empty()) {
    return nullptr;
  }
  if (is_stack) {
    auto iter = nodes.begin() + (level_end - 1);
    auto p = *iter;
    nodes.erase(iter);
    return p;
  } else {
    auto p = nodes.front();
    nodes.erase(nodes.begin());
    return p;
  }
}
inline void push_next(TreeNode *p, std::vector<TreeNode *> &nodes,
                      int level_end, int is_stack) {
  if (is_stack) {
    if (p->right) {
      nodes.insert(nodes.begin() + level_end, p->right);
    }
    if (p->left) {
      nodes.insert(nodes.begin() + level_end, p->left);
    }
  } else {
    if (p->left) {
      nodes.push_back(p->left);
    }
    if (p->right) {
      nodes.push_back(p->right);
    }
  }
}

std::vector<std::vector<int>> zigzagLevelOrder(TreeNode *root) {
  if (!root) {
    return {};
  }
  std::vector<std::vector<int>> ret{{}};
  std::vector<TreeNode *> level_cache{root};
  int level_cnt = 1;
  int is_stack = 0;
  while (!level_cache.empty()) {
    auto p = pop_next(level_cache, level_cnt, is_stack);
    ret.back().push_back(p->val);
    --level_cnt;

    push_next(p, level_cache, level_cnt, is_stack);
    if (level_cnt == 0) {
      level_cnt = level_cache.size();
      is_stack = (is_stack + 1) % 2;
      if (level_cnt) {
        ret.push_back({});
      }
    }
  }
  return ret;
}

int main() {
  auto root = new TreeNode(3, new TreeNode(9),
                           new TreeNode(20, new TreeNode(15), new TreeNode(7)));
  auto ret = zigzagLevelOrder(root);
  for (const auto v : ret) {
    PrintVect(v);
  }
}