20-Coding-Patterns-DesignGurus.org/8_binary_tree_level_order_traversal.cpp

/* Problem: https://leetcode.com/problems/binary-tree-level-order-traversal/
 */

#include <algorithm>
#include <climits>
#include <cstddef>
#include <iostream>
#include <limits>
#include <numeric>
#include <ostream>
#include <queue>
#include <vector>

//  Below TreeNode definition is copied from the problem
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) {}
};
// Above TreeNode definition is copied from the problem

std::ostream &operator<<(std::ostream &out, TreeNode *node) {
  std::queue<TreeNode *> bfs;
  if (node)
    bfs.push(node);
  out << "{";
  while (bfs.size()) {
    node = bfs.front();
    bfs.pop();
    if (node) {
      out << node->val << ", ";
      bfs.push(node->left);
      bfs.push(node->right);
    } else
      out << "NULL, ";
  }
  out << "}";
  return out;
}

std::ostream &operator<<(std::ostream &out, const std::vector<int> &vec) {
  out << "{";
  for (int num : vec) {
    std::cout << num << ", ";
  }
  out << " }";
  return out;
}

std::ostream &operator<<(std::ostream &out,
                         const std::vector<std::vector<int>> &vec_2d) {
  out << "{";
  for (const std::vector<int> &vec : vec_2d)
    out << vec << ", ";
  out << " }";
  return out;
}

TreeNode *vector_to_tree(std::vector<int> &&vec) {
  if (!vec.size())
    return nullptr;

  std::vector<TreeNode *> *nodes = new std::vector<TreeNode *>;
  for (std::size_t i = 0, upto = vec.size(); i < upto; i++) {
    if (vec[i] == INT_MIN)
      nodes->emplace_back(nullptr);
    else
      nodes->emplace_back(new TreeNode(vec[i]));
  }

  /* If input vector is a malformed tree with missing intermediate nodes than
   * below loop will cause segfault*/
  for (std::size_t i = 0, upto = vec.size(); i < upto; i++) {
    std::size_t left = 2 * i + 1, right = 2 * i + 2;
    if (left < upto)
      (*nodes)[i]->left = (*nodes)[left];
    if (right < upto)
      (*nodes)[i]->right = (*nodes)[right];
  }

  return (*nodes)[0];
}

std::vector<std::vector<int>> solution(TreeNode *node) {
  std::vector<std::vector<int>> tree_by_level;
  std::queue<TreeNode *> current_q;
  if (node != nullptr) {
    current_q.push(node);
  }
  while (current_q.size()) {
    std::queue<TreeNode *> next_q;
    std::vector<int> current_level;
    while (current_q.size()) {
      TreeNode *current_node = current_q.front();
      current_q.pop();
      current_level.push_back(current_node->val);
      if (current_node->left)
        next_q.push(current_node->left);
      if (current_node->right)
        next_q.push(current_node->right);
    }
    tree_by_level.push_back(current_level);
    current_q = next_q;
  }
  return tree_by_level;
}

int main() {
  std::cout << vector_to_tree({3, 9, 20, INT_MIN, INT_MIN, 15, 17}) << '\n';
  std::cout << solution(vector_to_tree({3, 9, 20, INT_MIN, INT_MIN, 15, 17}))
            << '\n'
            << solution(vector_to_tree({1})) << '\n'
            << solution(vector_to_tree({})) << '\n';
  return 0;
}