653. Two Sum IV - Input is a BST

はじめに

653.Two Sum IV - Input is a BST です。結構苦戦しました。1 時間はかかった気がする。 まずは二分木、加えて二分探索木というものを扱うすべをあまり知らなかったところが大きい。 もちろん、学生時代に木構造は学んでいたものの…思い出せない…

問題を理解するのは簡単です。二分探索木とある値 k が与えられるので、二分探索木のある 2 つのノードの値の合計がある値 k になるかを調べ、なるなら true、ならないなら false を返す、という問題です。 この問題は 1.Two Sum の類似問題として引用されていました。 なので、解き方は Two Sum を知っていれば少し参考になる部分があります。 それでは辛い初手をご覧ください。

初手(深さ優先探索:DFS)

ひとまず、あまりにも詰まったので途中で Grok 3 に相談してます。そのとき、

  • root ノードの値を map に入れ忘れていること
  • グローバル変数を使うと、テストケース全体で使い回され、予期せぬ挙動をすること の 2 点を指摘され、直しました。

そしてこの 2 点を改善し AC となったのが以下のコードです。
https://leetcode.com/problems/two-sum-iv-input-is-a-bst/submissions/1610648203

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func findTarget(root *TreeNode, k int) bool {
    hm := make(map[int]int)
    index := 0
    hm[root.Val] = index
    index = search(root.Left, hm, index)
    _ = search(root.Right, hm, index)
    for val, i := range hm {
        c := k - val
        if j, ok := hm[c]; ok && i != j {
            return true
        }
    }
    return false
}

func search(n *TreeNode, hm map[int]int, index int) int {
    if n == nil {
        return index
    }
    index += 1
    hm[n.Val] = index
    if n.Left != nil {
        index = search(n.Left, hm, index)
    }
    if n.Right != nil {
        index = search(n.Right, hm, index)
    }
    return index
}

僕の大の苦手な再帰関数です。どうにか自力で書きました。 やっていることは簡単です。単に与えられた根のノードの左のノードと右のノードを search という再帰関数に渡し、その再帰関数は map に値を入れつつ、インデックスを増加させていきます。 1.Two Sum でやった通り、すべての値とその値の添字さえあれば単一のループで解を導けます。 ノードのすべての値とその添字を持つ map である hm をループし、Two Sum よろしく k ー value を計算し、その値が map にあるか確認します。このとき、同じ値を利用しないよう注意しつつ、値があれば true、なければ false を返せばいいだけです。 計算量は O(n) です。

解けたとはいえ、すごくゴリ押しした感が否めません。計算量も悪くはないはずですが、スマートさを微塵も感じません。Solutions を見て改善します。

再帰しない DFS

AtCoder をしていたときに両方書いていた記憶はありますが、DFS は必ずしも再帰関数を用いなくとも解ける・・・はずです。 なので、単一の関数内のループで実装してみました。
https://leetcode.com/problems/two-sum-iv-input-is-a-bst/submissions/1611115653

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func findTarget(root *TreeNode, k int) bool {
    stack := []*TreeNode{root}
    hm := make(map[int]int)
    hm[root.Val] = 0
    index := 1
    for len(stack) > 0 {
        node := stack[len(stack)-1]
        stack = stack[:len(stack)-1]
        hm[node.Val] = index
        index += 1
        if node.Left != nil {
            stack = append(stack, node.Left)
        }
        if node.Right != nil {
            stack = append(stack, node.Right)
        }
    }
    for val, i := range hm {
        c := k - val
        if j, ok := hm[c]; ok && i != j {
            return true
        }
    }
    return false
}

スタックが肝です。DFS には変わりありません。BFS(深さ優先探索)にするなら、キューを使います。 ただ、上記のコードで少し気持ち悪いのは、ルートノードの右側から探索が始まる点です。 スタックは LIFO なので、次のループで取り出されるノードは、最後に積んだノードです。 上記のコードだと、以下のことです。

1
2
3
        if node.Right != nil {
            stack = append(stack, node.Right)
        }

なので、例1だと、⑥のノードから探索が始まります。 example1

ところで Solution をぱっと見たところ、インデックスの管理なんてしていませんでした。つまり、Two Sum に引きずられて index が重複していないかを見る必要なんてないようです。 次はその解法を試します。

最終手

もう solutions 見ましたよね。DFS で解いています。以前よりシンプルなコードになったのがおわかりでしょうか。 1.Two Sum に引きずられてインデックスの重複を避ける考え方を省きました。 単に、探索中に出てくる数字を記憶していき、k - n が記憶した数字の中にあれば true を返す、というものです。 再帰代表みたいな書き方ですね。知らんけど。
https://leetcode.com/problems/two-sum-iv-input-is-a-bst/submissions/1613601062

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func findTarget(root *TreeNode, k int) bool {
    return dfs(root, k, map[int]bool{})
}

func dfs(node *TreeNode, k int, m map[int]bool) bool {
    if node == nil {
        return false
    }
    if m[k - node.Val] {
        return true
    }
    m[node.Val] = true
    return dfs(node.Left, k, m) || dfs(node.Right, k, m)
}

計算量は、再帰によるノードの探索とハッシュマップの探索がノード数分かかりますが、単に O(n) で良いと思います。 LeetCode の AI もそう言っています。 time complexity

なお、この解法は Beats 100.00% 👏 です。

Hugo で構築されています。
テーマ StackJimmy によって設計されています。