Python程式:查詢二叉樹中長度為k的路徑


假設我們有一個包含唯一值的二叉樹,還有一個值k,我們需要找到樹中k長度的唯一路徑的數量。路徑可以從父節點到子節點,也可以從子節點到父節點。當一條路徑中出現某個節點而另一條路徑中沒有出現時,我們將認為這兩條路徑不同。

因此,如果輸入如下所示:

k = 3,則輸出為4,因為路徑為[12,8,3],[12,8,10],[8,12,15],[3,8,10]。

為了解決這個問題,我們將遵循以下步驟:

  • 定義一個函式dfs()。這將接收節點作為引數。

    • 如果節點為空,則

      • 返回一個列表,其中包含1個1和k-1個0。

    • left := dfs(節點的左子節點)

    • right := dfs(節點的右子節點)

    • 對於i從0到K,執行以下操作:

      • ans := ans + left[i] * right[K - 1 - i]

    • res := 一個大小為K的包含0的列表

    • res[0] := 1, res[1] := 1

    • 對於i從1到K - 1,執行以下操作:

      • res[i + 1] := res[i + 1] + left[i]

      • res[i + 1] := res[i + 1] + right[i]

    • 返回res

  • 在主方法中,執行以下操作:

  • ans := 0


  • dfs(根節點)


  • 返回ans


讓我們看看下面的實現,以便更好地理解:

示例

線上演示

class TreeNode:
   def __init__(self, data, left = None, right = None):
      self.data = data
      self.left = left
      self.right = right
class Solution:
   def solve(self, root, K):
      def dfs(node):
         if not node:
            return [1] + [0] * (K-1)
         left = dfs(node.left)
         right = dfs(node.right)
         for i in range(K):
            self.ans += left[i] * right[K - 1 - i]
         res = [0] * K
         res[0] = res[1] = 1
         for i in range(1, K - 1):
            res[i + 1] += left[i]
            res[i + 1] += right[i]
         return res
      self.ans = 0
      dfs(root)
      return self.ans
ob = Solution()
root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)
print(ob.solve(root, 3))

輸入

root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)
3

輸出

4

更新於:2020年10月6日

瀏覽量:227

開啟你的職業生涯

完成課程獲得認證

開始學習
廣告