LeetCode/C++/637.cpp

#include <iostream>
#include <cstdint>
#include <vector>

#include "./utils.hpp"

using std::vector;
using std::string;

/**
 * Definition for a binary tree node.
 * 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) {}
 * };
 */
class Solution {
public:
    vector<double>
    averageOfLevels(TreeNode *root)
    {
        vector<double> vec;
        vector<int> amount;
        _go(0, root, vec, amount);
        for ( unsigned int i = 0; i < vec.size(); ++i ) {
            vec[i] /= amount[i];
        }
        return vec;
    }

    void
    _go(unsigned long level, TreeNode *root, vector<double> &vec, vector<int> &amount)
    {
        if ( root == NULL ) { return; }
        if ( level >= vec.size() ) {
            vec.push_back(0);
            amount.push_back(0);
        }
        vec[level] += root->val;
        ++amount[level];
        ++level;
        _go(level, root->left, vec, amount);
        _go(level, root->right, vec, amount);
    }
};


void r(TreeNode *root, vector<double> exp) {
    std::cout << "Solution().averageOfLevels(" << treenode_to_str(root) << ") = ";
    std::cout << vector_to_str(Solution().averageOfLevels(root));
    std::cout << " | exp: " << vector_to_str(exp) << "\n";
};

int main(void) {
    TreeNode root1l = TreeNode(9);
    TreeNode root1rl = TreeNode(15);
    TreeNode root1rr = TreeNode(7);
    TreeNode root1r = TreeNode(20, &root1rl, &root1rr);
    TreeNode root1 = TreeNode(3, &root1l, &root1r);
    r(&root1, {3.00000,14.50000,11.00000});

    TreeNode root2ll = TreeNode(15);
    TreeNode root2lr = TreeNode(7);
    TreeNode root2l = TreeNode(9, &root2ll, &root2lr);
    TreeNode root2r = TreeNode(20);
    TreeNode root2 = TreeNode(3, &root2l, &root2r);
    r(&root2, {3.00000,14.50000,11.00000});

    return 0;
}