オットセイの経営日誌

データサイエンス系ベンチャーを経営してます。経営のこと、趣味のことつぶやきます。

LeetCode / Path Sum II

https://leetcode.com/problems/path-sum-ii/

Given a binary tree and a sum, find all root-to-leaf paths where each path's sum equals the given sum.

Note: A leaf is a node with no children.

Example:

Given the below binary tree and sum = 22,

f:id:mhiro216:20190807002733p:plain:w100

Return:

f:id:mhiro216:20190807002746p:plain:w100

前回記事の発展系で、
binary treeの先端から末端までのroot-to-leafの合計値が、変数sumの値と合致するpathを全て返す問題です。

解答・解説

解法の説明に入る前に、本問題のように経路を順に辿ってゴールするような問題の解法には、

の2種類があるそうです。
各々の意味と具体的な解法を示します。

解法1

私がsubmitしたIterativeな手法です。久々に一発でsubmitが通って嬉しかったです。
こちらがDFSにあたります。

treeの各rootを巡回しながら、3つの要素から成るtupleを変数stackに格納します。
1つ目の要素は現在いるroot、2つ目の要素は辿ったrootの値の合計、3つ目の要素は辿ったrootの値のリストです。
stack.pop()で、後に格納したtupleから取り出します。
取り出したtupleの2つ目の要素にrootの値を足し(と同時に3つ目の要素にrootの値をappendし)、値がsumと同値になった時点の3つ目の要素を、最終的な回答となる変数ansに格納します。

stackの後方に格納されているtupleほどtreeの深い階層にある要素になるので、stack.pop()で後に格納したtupleから取りしてwhile loopにかけるということは、より深い階層から先に探索している(=深さ優先探索をしている)ことになります。

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def pathSum(self, root: TreeNode, sum: int) -> List[List[int]]:
        ans = []
        if not root:
            return ans
        stack = [(root, root.val, [root.val])]
        while stack:
            curr, val, path = stack.pop()
            if not curr.left and not curr.right:
                if val == sum:
                    ans.append(path)
            if curr.right:
                stack.append((curr.right, val+curr.right.val, path+[curr.right.val]))
            if curr.left:
                stack.append((curr.left, val+curr.left.val, path+[curr.left.val]))
        return ans
解法2

続いてIterativeなBFS。
treeの各rootを巡回しながら、3つの要素から成るtupleを変数queueに格納します。
1つ目の要素は現在いるroot、2つ目の要素は辿ったrootの値の合計、3つ目の要素は辿ったrootの値のリストです。ここまではDFSと同じ。
DFSと異なるのは、queue.pop(0)で、先に格納したtupleから取り出すところです。

queueの前方に格納されているtupleほどtreeの浅い階層にある要素になるので、queue.pop(0)で前に格納したtupleから取りしてwhile loopにかけるということは、より浅い階層から先に探索している(=幅優先探索をしている)ことになります。

class Solution:
    def pathSum(self, root: TreeNode, sum: int) -> List[List[int]]:
        ans = []
        if not root:
            return ans
        queue = [(root, root.val, [root.val])]
        while queue:
            curr, val, path = queue.pop(0)
            if not curr.left and not curr.right:
                if val == sum:
                    ans.append(path)
            if curr.right:
                queue.append((curr.right, val+curr.right.val, path+[curr.right.val]))
            if curr.left:
                queue.append((curr.left, val+curr.left.val, path+[curr.left.val]))
        return ans
解法3

最後にRecursiveな解法は以下の通りです。考え方は前の解法に近しいので解説は省略します。
しかしなかなかRecursiveな解法を自力で完成させられません。脳がIterativeな考え方に慣れきってますね…

class Solution:
    def pathSum(self, root: TreeNode, sum: int) -> List[List[int]]:
        if not root:
            return []
        res = []
        self.dfs(root, sum, [], res)
        return res
    def dfs(self, root, sum, ls, res):
        if not root.left and not root.right and sum == root.val:
            ls.append(root.val)
            res.append(ls)
        if root.left:
            self.dfs(root.left, sum-root.val, ls+[root.val], res)
        if root.right:
            self.dfs(root.right, sum-root.val, ls+[root.val], res)