/* Problem: https://leetcode.com/problems/path-sum-iii/
 */

#include "lib_leetcode.h"
#include <climits>
#include <iostream>
#include <stack>
#include <vector>

int solution(TreeNode *root, const int target_sum) {
  if (!root)
    return 0;

  int sum_paths = 0;
  std::vector<long long> paths_sum;
  std::stack<node_and_explored> dfs;
  dfs.push({root, false});
  while (dfs.size()) {
    auto &[node, explored] = dfs.top();
    int node_val = node->val;
    if (explored) {
      dfs.pop();
      paths_sum.pop_back();
      for (long long &sum : paths_sum)
        sum -= node_val;
    } else {
      explored = true;
      paths_sum.push_back(0);
      for (long long &sum : paths_sum) {
        sum += node_val;
        if (sum == target_sum)
          sum_paths++;
      }
      TreeNode *left_child = node->left, *right_child = node->right;
      if (right_child)
        dfs.push({right_child, false});
      if (left_child)
        dfs.push({left_child, false});
    }
  }

  return sum_paths;
}

int main() {
  std::cout << solution(vector_to_tree(
                            {10, 5, -3, 3, 2, INT_MAX, 11, 3, -2, INT_MAX, 1}),
                        8)
            << '\n'
            << solution(vector_to_tree({5, 4, 8, 11, INT_MAX, 13, 4, 7, 2,
                                        INT_MAX, INT_MAX, 5, 1}),
                        22)
            << '\n';
}