ZigZag binary tree traversal in C++

Question

I am trying to attempt zig zag traversal of binary tree. But I am stuck at one type of test cases, ie, when the tree is not balanced. If I give my input as

3
3 9 20 null null 15 7

for a binary tree which looks like this:

    3
   / \
  9  20
    /  \
   15   7

I get output:

3 
20 9 
0 

If my input was 3 9 20 1 null 15 7, my output is: 3 20 9 1 0

Below is my code.

struct node
    {
        int data;
        node* left;
        node* right;
    };

void order (node* root, map <int, vector <int> >& ans, int level, int k) {
    if (root == NULL) {
        return;
    }

    if (k==0) {
        if (root->left != NULL) {
            ans[level+1].push_back(root->left->data);
            order (root->left, ans, level+1, 1);
        }
        if (root->right != NULL) {
            ans[level+1].push_back(root->right->data);
            order (root->right, ans, level+1, 1);
        }
    }
    else if (k==1) {
        order (root->left, ans, level+1, 0);
        order (root->right, ans, level+1, 0);

        if (root->right != NULL)
            ans[level+1].push_back(root->right->data);
        if (root->left != NULL)
            ans[level+1].push_back(root->left->data);
    }
}

vector<vector<int> > zigzag(node* root){
     map <int, vector <int> > ans;
     vector <vector <int> > zig;

     ans[0].push_back(root->data);

     order(root, ans, 1, 1);

     for (auto it = ans.begin(); it != ans.end(); it++ ) {   //converting map into vector
        zig.push_back(it->second);
     }
     return zig;
 }

---

From what I understand, if the input is null my code should return nothing and continue execution of further nodes. I can't figure out my mistake. Can anybody help me? TIA!

Answer 1

Your code actually appears to work with your provided test case. I suspect your problem there is with input/output. However the solution itself is unfortunately not. Consider the following test case:

      1
     / \
    5   2
   /     \
  6       3
 /         \
7           4

Your solution will output: [[1], [5, 2], [6, 3], [7, 4]]

Let's call each vector in your zigzag vector a level.

1. First you add the 1 (the root) to your first level.
2. k==1 level==1 Then you recurse left.
3. k==0 level==2 You add 6 correctly to level 2. And recurse left again.
4. k==1 level==3 Can't recurse left or right. Add 7 incorrectly to level 3. Return
5. k==0 level==2 Can't recurse right. Return.
6. k==1 level==1 Add 5 incorrectly to level 1. Recurse right.
7. etc.

I think this approach is going to be difficult to get to the correct solution. Primarily because of it's DFS nature. A BFS solution could be the right path. It is naturally suited to these level-by-level styled problems.

Answer 2

For level order traversal you can use the stack data structure for traversing the each level in zig zag fashion.

vector<vector<int> > Solution::zigzagLevelOrder(TreeNode* root) {
    stack<TreeNode*> s1, s2;
    s1.push(root);
    vector< vector< int> > ans;
    int check = 1;
    while(!s1.empty() or !s2.empty()){
        vector<int> current;
        if(check){
            while(!s1.empty()){
                root = s1.top();
                s1.pop();
                current.push_back(root->val);
                if(root->left)
                    s2.push(root->left);
                if(root->right)
                    s2.push(root->right);
            }
            check = 1 - check;
        }
        else{
            while(!s2.empty()){
                root = s2.top();
                s2.pop();
                current.push_back(root->val);
                if(root->right)
                    s1.push(root->right);
                if(root->left)
                    s1.push(root->left);
            }
            check = 1 - check;
        }

        ans.push_back(current);
    }
    return ans;

}

Maintain the first stack for odd level and second stack for even level. While traversing the odd level push left child then right, and while traversing even level traverse right child then left.