Swift で LCA: Lowest Common Ancestor（最近共通祖先)

備忘メモ

NOTE

• 実装：dfs (post-order)
• 問題：https://leetcode.com/problems/lowest-common-ancestor-of-a-binary-tree/

コード

func lowestCommonAncestor(_ root: TreeNode?, _ p: TreeNode?, _ q: TreeNode?) -> TreeNode? {
  // Return null if not find nodes we are looking for.
  func dfs(_ root: TreeNode?) -> TreeNode? {
    if root == nil { return nil }
    if root?.val == p?.val || root?.val == q?.val { return root } // current is node we are looking for.
    let left  = dfs(root?.left)
    let right = dfs(root?.right)
    if left == nil && right == nil { return nil } // if not find node in children.
    if left != nil && right != nil { return root } // if current is LCA.
    return left == nil ? right : left // return either if one child found node.
  }
  
  return dfs(root)
}

ポイント

dfs がreturnする TreeNode だが、再帰的に呼び出されてる中で途中で意味合いがことなる。それを１つのTreeNodeという表現かつロジックで書いてあるのでやや混乱した。

再帰的に呼び出されてる dfs がreturn する値の意味は下記の２つのどちらかのケースになる。

1. 探し当てた p または q の位置
2. LCA の位置

1. の通りp か q のどちらか１つを見つけたなら return するのはその見つけた p か q になるが、2) のようにchildren の両方ともが p と q を見つけてきてる場合はこのノードがLCAになる。なので、そこから上へ渡すのはLCAになる。