55 lines
1.3 KiB
C++
55 lines
1.3 KiB
C++
/* 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';
|
|
}
|