leetcode/count_unival_subtrees.cpp at master · Mayank-Parasar/leetcode · 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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
//
// Created by Mayank Parasar on 2020-01-11.
//

/*
 * A unival tree is a tree where all the nodes have the same value. Given a binary tree,
 * return the number of unival subtrees in the tree.

For example, the following tree should return 5:

   0
  / \
 1   0
    / \
   1   0
  / \
 1   1

The 5 trees are:
- The three single '1' leaf nodes. (+3)
- The single '0' leaf node. (+1)
- The [1, 1, 1] tree at the bottom. (+1)
 */

#include <iostream>
#include <vector>
#include <assert.h>

using namespace std;

struct node {
    int val;

    node* left;
    node* right;

    node(int x) : val(x), left(nullptr), right(nullptr)
    {}
};

bool is_unival(node* node_) {

    if (node_->left == nullptr &&
        node_->right == nullptr)
        return true;

    if (is_unival(node_->left) && is_unival(node_->right) &&
        node_->val == node_->left->val &&
        node_->val == node_->right->val)
        return true;

    if(node_->left == nullptr) {
        if(is_unival(node_->right) && node_->val == node_->right->val)
            return true;
    }

    if(node_->right == nullptr) {
        if(is_unival(node_->left) && node_->val == node_->left->val)
            return true;
    }

}

int total_count = 0;

int count_univals(node* node_) {
    if(node_ == nullptr)
        return 0;

    total_count = count_univals(node_->left) + count_univals(node_->right);

    if(is_unival(node_))
        total_count +=1;

    return total_count;
}

int main() {

    node* node1 = new node(0);
    node* node2 = new node(1);
    node* node3 = new node(0);
    node* node4 = new node(1);
    node* node5 = new node(0);
    node* node6 = new node(1);
    node* node7 = new node(1);

    // connect them

    node1->left = node2;
    node1->right = node3;
    node3->left = node4;
    node3->right = node5;
    node4->left = node6;
    node4->right = node7;

    cout << boolalpha;

    cout << is_unival(node2);
    cout << endl;
    cout << "number of unival trees present: " << count_univals(node1);


    return 0;
}