URL

Split Linked List in Parts - LeetCode

Code

https://github.com/syohex/dotnet-study/tree/master/fsharp/leetcode/problems/0725/main.fsx

type ListNode =
    | Leaf
    | Node of int * ListNode

let listLength (head: ListNode) : int =
    let rec listLength' node acc =
        match node with
        | Leaf -> acc
        | Node(_, next) -> listLength' next (acc + 1)

    listLength' head 0

let rec listToListNode ns =
    match ns with
    | [] -> Leaf
    | h :: t -> Node(h, listToListNode t)

let splitListToParts (head: ListNode) (k: int) : ListNode list =
    let rec splitListToParts' node count elements mods acc ret =
        match node with
        | Leaf ->
            match acc with
            | [] -> ret
            | _ -> (listToListNode (List.rev acc)) :: ret
        | Node(i, next) ->
            let acc' = i :: acc

            if count - 1 = 0 then
                let count' = if mods > 0 then elements + 1 else elements
                splitListToParts' next count' elements (mods - 1) [] ((listToListNode (List.rev acc')) :: ret)
            else
                splitListToParts' next (count - 1) elements mods acc' ret

    let len = listLength head
    let elements = len / k
    let mods = if len % k = 0 then 0 else len % k
    let nulls = if len < k then k - len else 0
    let count = if mods > 0 then elements + 1 else elements

    let splited = splitListToParts' head count elements (mods - 1) [] []
    seq { 1..nulls } |> Seq.fold (fun acc _ -> Leaf :: acc) splited |> List.rev