F# Data Structures / DataStructures / AvlTree.fs

module AvlTree =

    type 'a AvlTree = 
    | Empty
    | Node of (int * 'a) * 'a AvlTree * 'a AvlTree


    let inline getHeight t = 
        match t with
        | Empty -> 0
        | Node((h, _), _, _) -> h


    let inline isBalanced node = 
        match node with
        | Empty -> true
        | Node(_, l, r) -> abs(getHeight l - getHeight r) <= 1


    //Calculates height of a node based on the heights of it's children
    let inline newHeight l r = 
        let lH = getHeight l
        let rH = getHeight r
        (max lH rH) + 1


    let rec contains tree value = 
        match tree with
        | Empty -> false
        | Node((_, v), l, r) -> 
            let c = compare value v
            if c > 0 
            then contains r value
            elif c < 0 
            then contains l value
            else true


    let restructure parent = 

        match parent with
        | Empty -> failwith "parent node cannot be empty when checking if it needs restructuring"
        | Node((h, v), l, r) ->
        
            let fixR rv rl rr =
                let inline useRR () = 
                    let l'= Node((newHeight l rl, v), l, rl)
                    let r'= rr
                    Node((newHeight l' r', rv), l', r')
                    
                let inline useRL rlv rll rlr = 
                    let l'= Node((newHeight l rll, v), l, rll)
                    let r'= Node((newHeight rlr rr, rv), rlr, rr)
                    Node((newHeight l' r', rlv), l', r')
                    
                match rl, rr with
                | Empty, Empty -> failwith (sprintf "both grand children cannot be empty of child node %A" r)
                | Empty, Node((rrh, _), _, _) -> assert (rrh > 0); useRR ()
                | Node((rlh, rlv), rll, rlr), Empty -> assert (rlh > 0); useRL rlv rll rlr
                | Node((rlh, _), _, _), Node((rrh, rrv), rrl, rrr) when rrh >= rlh -> useRR ()
                | Node((_, rlv), rll, rlr), Node(_) -> useRL rlv rll rlr


            let fixL lv ll lr = 
                let inline useLR lrv lrl lrr = 
                    let l'= Node((newHeight ll lrl, lv), ll, lrl)
                    let r'= Node((newHeight lrr r, v), lrr, r)
                    Node((newHeight l' r', lrv), l', r')

                let inline useLL () = 
                    let l'= ll
                    let r'= Node((newHeight lr r, v), lr, r)
                    Node((newHeight l' r', lv), l', r')
                     
                match ll, lr with
                | Empty, Empty -> failwith (sprintf "both grand children cannot be empty of child node %A" l)
                | Empty, Node((lrh, lrv), lrl, lrr) -> assert (lrh > 0); useLR lrv lrl lrr
                | Node((llh, _), _, _), Empty -> assert(llh > 0); useLL ()
                | Node((llh, _), _, _), Node((lrh, _), _, _) when llh >= lrh -> useLL ()
                | Node(_), Node((_, lrv), lrl, lrr) -> useLR lrv lrl lrr
            
            match l, r with
            | Empty, Empty -> parent
            | Empty, Node((1, _), _, _) -> parent
            | Node((1, _), _, _), Empty -> parent
            | Node((lh, _), _, _), Node((rh, _), _, _) when abs(lh - rh) <= 1 -> parent
            | Empty, Node((rh, rv), rl, rr) -> assert (rh > 1); fixR rv rl rr
            | Node((lh, lv), ll, lr), Empty -> assert (lh > 1); fixL lv ll lr
            | Node((lh, lv), ll, lr), Node((rh, _), _, _) when lh > rh -> fixL lv ll lr
            | Node(_), Node((_, rv), rl, rr) -> fixR rv rl rr



    //if the value is already present in the tree when 
    //we do insertion, then just return the original tree

    //Maybe we need to ovewrite the old value with the new one ?
    let insert tree value = 
        let rec doInsert subTree cont =
            match subTree with
            | Empty -> cont(Node((1, value), Empty, Empty))
            | Node((h, v), lc, rc) -> 
                let c = compare value v
                if c > 0 
                then doInsert rc (fun rc' ->
                        let parent' = Node((newHeight lc rc', v), lc, rc')
                        cont(restructure parent'))
                elif c < 0 
                then doInsert lc (fun lc' -> 
                        let parent' = Node((newHeight lc' rc, v), lc', rc)
                        cont(restructure parent'))
                else tree
        doInsert tree (fun n -> n)





    let remove tree value = 
        
        let rec doRemove subTree cont =
            match subTree with 
            | Empty -> failwith "subtree cannot be empty when removing, \
                        since that was handled before entering this function"
            | Node((h, v), Empty, r) -> cont(v, r)
            | Node((h, v), l, r) -> doRemove l (fun (u, l') -> 
                                        let n' = Node((newHeight l' r, v), l', r)
                                        cont(u, restructure n'))


        let rec findNodeAndRemove subTree cont =
            match subTree with
            | Empty -> tree 
            | Node((h, v), l, r) ->
                let c = compare value v
                if c > 0
                then findNodeAndRemove r (fun r' -> 
                    let newNode = Node((newHeight l r', v), l, r')
                    cont(restructure newNode))
                elif c < 0 
                then findNodeAndRemove l (fun l' -> 
                    let n' = Node((newHeight l' r, v), l', r )
                    cont(restructure n'))
                else 
                    match l, r with
                    | Empty, Empty -> cont(Empty)
                    | _, Empty -> cont(l)
                    | _, _ -> 
                        let u, r' = doRemove r (fun (v, n) -> v, n)
                        let n' = Node((newHeight l r', u), l, r')
                        cont(restructure n')

        findNodeAndRemove tree (fun n -> n)
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.