在二叉树中找到累加和为指定值的最长路径长度

在二叉树中找到累加和为指定值的最长路径长度

题目

给定一棵二叉树根节点root和一个正数k，二叉树节点值类型为整型，求累加和为k的最长路径。路径是指从根节点到任意一个节点的路径的子路径。

例如

• k=6，可以找到从根的一条路径 [-3, 3, 0, 6]
• k=-9，可以找到从根的一条路径-3, [-9]，子路径-9满足和为-9.
• k=3，可以找到从根的一条路径-3, [3,0]，子路径[3, 0]满足和为3

思路

类似未排序整数数组中累加和为给定值的最长子数组长度
先序遍历，从根节点到当前节点的路径理解成链接所示题目的数组，则该路径的子数组为所求。
需要：

• 记录遍历到节点所在树的层
• 从根到遍历到的节点累加得到的和-层映射关系

实现

def max_length_path_with_sum(root, k):
    def pre_order(node, pre_sum, level, max_len):
        if node is None:
            return max_len

        cur_sum = pre_sum + node.val
        if cur_sum not in sum_map:
            sum_map[cur_sum] = level

        if cur_sum-k in sum_map:
            max_len = max(max_len, level-sum_map[cur_sum-k])

        max_len = pre_order(node.left, cur_sum, level+1, max_len)
        max_len = pre_order(node.right, cur_sum, level+1, max_len)

        if sum_map[cur_sum] == level:
            sum_map.pop(cur_sum)

        return max_len

    sum_map = {0: 0}
    return pre_order(root, pre_sum=0, level=1, max_len=0)

测试

没有找到很好的测试方式，也没找到leetcode类似的题目，先用简单的例子测试

def test_max_length_path_with_sum():
    vals = [-3, 3, -9, 1, 0, 2, 1, 1, 6]
    nodes = [TreeNode(v) for v in vals]

    root = nodes[0]
    nodes[0].left, nodes[0].right = nodes[1], nodes[2]
    nodes[1].left, nodes[1].right = nodes[3], nodes[4]
    nodes[2].left, nodes[2].right = nodes[5], nodes[6]
    nodes[4].left, nodes[4].right = nodes[7], nodes[8]

    assert(max_length_path_with_sum(root, 6) == 4)
    assert(max_length_path_with_sum(root, -9) == 1)
    assert(max_length_path_with_sum(root, -10) == 3)
    assert(max_length_path_with_sum(root, -11) == 3)
    assert(max_length_path_with_sum(root, 4) == 3)
    assert(max_length_path_with_sum(root, 1) == 4)
    print(max_length_path_with_sum(root, 3))
    assert(max_length_path_with_sum(root, 3) == 3)


if __name__ == '__main__':
    test_max_length_path_with_sum()