算法题/树的DFS：一个和的路径数

题目

给定一棵二叉树和一个数字S，找出树中的所有路径，使每条路径的所有节点值的和等于S。请注意，路径可以在任何节点开始或结束，但所有路径必须遵循从父到子(从上到下)的方向。

解答

Java

import java.util.*;

class TreeNode {
  int val;
  TreeNode left;
  TreeNode right;

  TreeNode(int x) {
    val = x;
  }
};

class CountAllPathSum {
  public static int countPaths(TreeNode root, int S) {
    List<Integer> currentPath = new LinkedList<>();
    return countPathsRecursive(root, S, currentPath);
  }

  private static int countPathsRecursive(TreeNode currentNode, int S, List<Integer> currentPath) {
    if (currentNode == null)
      return 0;

    // 添加当前节点到路径中
    currentPath.add(currentNode.val);
    int pathCount = 0, pathSum = 0;
    // 在当前路径列表中找到所有子路径的和
    ListIterator<Integer> pathIterator = currentPath.listIterator(currentPath.size());
    while (pathIterator.hasPrevious()) {
      pathSum += pathIterator.previous();
      // 如果任何子路径和S相等，便增加路径数
      if (pathSum == S) {
        pathCount++;
      }
    }

    // 遍历左子树
    pathCount += countPathsRecursive(currentNode.left, S, currentPath);
    // 遍历右子树
    pathCount += countPathsRecursive(currentNode.right, S, currentPath);

    // 从要回溯的路径中删除当前节点 
    // 当我们向上递归调用堆栈时，我们需要删除当前节点。
    currentPath.remove(currentPath.size() - 1);

    return pathCount;
  }

  public static void main(String[] args) {
    TreeNode root = new TreeNode(12);
    root.left = new TreeNode(7);
    root.right = new TreeNode(1);
    root.left.left = new TreeNode(4);
    root.right.left = new TreeNode(10);
    root.right.right = new TreeNode(5);
    System.out.println("Tree has path: " + CountAllPathSum.countPaths(root, 11));
  }
}

JavaScript

class TreeNode {
  constructor(val) {
    this.val = val;
    this.left = null;
    this.right = null;
  }
}

class CountAllPathSum {
  static countPaths(root, S) {
    const currentPath = [];
    return CountAllPathSum.countPathsRecursive(root, S, currentPath);
  }

  static countPathsRecursive(currentNode, S, currentPath) {
    if (currentNode === null)
      return 0;

    currentPath.push(currentNode.val);
    let pathCount = 0;
    let pathSum = 0;

    for (let i = currentPath.length - 1; i >= 0; i--) {
      pathSum += currentPath[i];
      if (pathSum === S) {
        pathCount++;
      }
    }

    pathCount += CountAllPathSum.countPathsRecursive(currentNode.left, S, currentPath);
    pathCount += CountAllPathSum.countPathsRecursive(currentNode.right, S, currentPath);

    currentPath.pop();

    return pathCount;
  }

  static main() {
    const root = new TreeNode(12);
    root.left = new TreeNode(7);
    root.right = new TreeNode(1);
    root.left.left = new TreeNode(4);
    root.right.left = new TreeNode(10);
    root.right.right = new TreeNode(5);
    console.log("Tree has path: " + CountAllPathSum.countPaths(root, 11));
  }
}

CountAllPathSum.main();

Python

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


class CountAllPathSum:
    def countPaths(root, S):
        currentPath = []
        return CountAllPathSum.countPathsRecursive(root, S, currentPath)

    def countPathsRecursive(currentNode, S, currentPath):
        if currentNode is None:
            return 0
        
        currentPath.append(currentNode.val)
        pathCount = 0
        pathSum = 0

        for i in range(len(currentPath)-1, -1, -1):
            pathSum += currentPath[i]
            if pathSum == S:
                pathCount += 1

        pathCount += CountAllPathSum.countPathsRecursive(currentNode.left, S, currentPath)
        pathCount += CountAllPathSum.countPathsRecursive(currentNode.right, S, currentPath)

        currentPath.pop()

        return pathCount

    def main():
        root = TreeNode(12)
        root.left = TreeNode(7)
        root.right = TreeNode(1)
        root.left.left = TreeNode(4)
        root.right.left = TreeNode(10)
        root.right.right = TreeNode(5)
        print("Tree has path:", CountAllPathSum.countPaths(root, 11))


CountAllPathSum.main()

Go

package main

import "fmt"

type TreeNode struct {
	val   int
	left  *TreeNode
	right *TreeNode
}

type CountAllPathSum struct{}

func (c *CountAllPathSum) countPaths(root *TreeNode, S int) int {
	currentPath := []int{}
	return c.countPathsRecursive(root, S, currentPath)
}

func (c *CountAllPathSum) countPathsRecursive(currentNode *TreeNode, S int, currentPath []int) int {
	if currentNode == nil {
		return 0
	}

	currentPath = append(currentPath, currentNode.val)
	pathCount := 0
	pathSum := 0

	for i := len(currentPath) - 1; i >= 0; i-- {
		pathSum += currentPath[i]
		if pathSum == S {
			pathCount++
		}
	}

	pathCount += c.countPathsRecursive(currentNode.left, S, currentPath)
	pathCount += c.countPathsRecursive(currentNode.right, S, currentPath)

	currentPath = currentPath[:len(currentPath)-1]

	return pathCount
}

func main() {
	root := &TreeNode{val: 12}
	root.left = &TreeNode{val: 7}
	root.right = &TreeNode{val: 1}
	root.left.left = &TreeNode{val: 4}
	root.right.left = &TreeNode{val: 10}
	root.right.right = &TreeNode{val: 5}

	c := &CountAllPathSum{}
	fmt.Println("Tree has path:", c.countPaths(root, 11))
}