小红的推理树平衡链路

import sys
from collections import Counter

class TreeNode:
    def __init__(self, val = 0, left = None, right = None):
        self.val = val
        self.left = left
        self.right = right

sys.setrecursionlimit(200000)

def build_tree_from_list(lst):
    if not lst or lst[0] is None:
        return None
    root = TreeNode(lst[0])
    queue = [root]
    i = 1
    while queue and i < len(lst):
        node = queue.pop(0)
        if lst[i] is not None:
            node.left = TreeNode(lst[i])
            queue.append(node.left)
        i += 1
        if i < len(lst) and lst[i] is not None:
            node.right = TreeNode(lst[i])
            queue.append(node.right)
        i += 1
    return root

def count_balanced_paths(root):
    cnt = Counter()
    cnt[0] = 1
    ans = 0
    zero_cnt = 0

    def dfs(node, s):
        nonlocal ans, zero_cnt
        if node is None:
            return

        s += node.val
        if node.val == 0:
            zero_cnt += 1

        ans += cnt[s]
        cnt[s] += 1
        dfs(node.left, s)
        dfs(node.right, s)
        cnt[s] -= 1

    dfs(root, 0)
    return ans - zero_cnt

def main():
    n = int(sys.stdin.readline().strip())
    line = sys.stdin.readline().strip().split()
    tree = [int(x) if x != "None" else None for x in line]

    root = build_tree_from_list(tree)
    result = count_balanced_paths(root)
    print(result)

if __name__ == "__main__":
    main()

import java.io.*;
import java.util.*;

class TreeNode{
    long val;
    TreeNode left;
    TreeNode right;
    TreeNode(long val){
        this.val = val;
        this.left = null;
        this.right = null;
    }
}
public class Main {
    static long ans = 0;
    static long zeroCount = 0;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int n = Integer.parseInt(br.readLine());
        if(n == 0){
            System.out.println(0);
            return;
        }
        String[] arr = br.readLine().split(" ");
        // 读取数据

        // 利用队列，以及层序遍历结果，构建树
        Queue<TreeNode> queue = new LinkedList<>();

        TreeNode root = new TreeNode(Long.parseLong(arr[0]));
        queue.offer(root);
        int idx = 1;

        while(!queue.isEmpty() && idx < n){
            TreeNode curr = queue.poll();
            if(idx < n && !arr[idx].equals("None")){
                curr.left = new TreeNode(Long.parseLong(arr[idx]));
                queue.offer(curr.left);
            }
            idx++;
            if(idx < n && !arr[idx].equals("None")){
                curr.right = new TreeNode(Long.parseLong(arr[idx]));
                queue.offer(curr.right);
            }
            idx++;
        }

        // 树建好了

        // 收集所有非空node，这些非null节点，可以作为遍历的起点
        List<TreeNode> allNodes = new ArrayList<>();
        Queue<TreeNode> nodeQueue = new LinkedList<>();
        nodeQueue.offer(root);

        while(!nodeQueue.isEmpty()){
            TreeNode node = nodeQueue.poll();
            allNodes.add(node);
            if(node.left != null) nodeQueue.offer(node.left);
            if(node.right != null) nodeQueue.offer(node.right);
        }


        Map<Long, Integer> map = new HashMap<>();
        map.put(0L, 1); // 初始化，前缀和0出现1次
        dfs(root, 0L, map);


        System.out.println(ans - zeroCount);
    }

    private static void dfs(TreeNode node, long sum, Map<Long, Integer> map){
        if(node == null) return;

        sum += node.val;

        if(node.val == 0){
            zeroCount++;
        }
        ans += map.getOrDefault(sum, 0);

        map.put(sum, map.getOrDefault(sum, 0) + 1);

        dfs(node.left, sum, map);
        dfs(node.right, sum, map);

        map.put(sum, map.get(sum) - 1);
    }
}