LeetCode 2458. Height of Binary Tree After Subtree Removal Queries in F#

URL

https://leetcode.com/problems/height-of-binary-tree-after-subtree-removal-queries/description/?envType=daily-question&envId=2024-10-26

Code

https://github.com/syohex/dotnet-study/blob/master/fsharp/leetcode/challenge/202410/height_of_binary_tree_after_subtree_removal_queries/main.fsx

type Tree =
    | Leaf
    | Node of int * Tree * Tree

let treeNodeHeights (root: Tree) : Map<int, int> =
    let rec treeNodeHeights' node acc =
        match node with
        | Leaf -> -1, acc
        | Node(v, left, right) ->
            let leftHeight, acc = treeNodeHeights' left acc
            let rightHeight, acc = treeNodeHeights' right acc
            let height = 1 + max leftHeight rightHeight
            height, Map.add v height acc

    treeNodeHeights' root Map.empty |> snd

let treeQueries (root: Tree) (queries: int list) : int list =
    let rec heightRemoveNodes' node depth maxHeight treeHeights acc =
        match node with
        | Leaf -> acc
        | Node(v, left, right) ->
            let leftMaxHeight =
                match left with
                | Leaf -> depth
                | Node(v, _, _) -> depth + 1 + Map.find v treeHeights

            let rightMaxHeight =
                match right with
                | Leaf -> depth
                | Node(v, _, _) -> depth + 1 + Map.find v treeHeights

            let acc = Map.add v maxHeight acc

            let acc =
                heightRemoveNodes' left (depth + 1) (max maxHeight rightMaxHeight) treeHeights acc

            heightRemoveNodes' right (depth + 1) (max maxHeight leftMaxHeight) treeHeights acc

    let treeHeights = treeNodeHeights root
    let heights = heightRemoveNodes' root 0 0 treeHeights Map.empty

    queries |> List.map (fun query -> Map.find query heights)