Simplify the Algebraic Expressions - https://www.hackerrank.com/contests/lambda-calculi-march-2016/challenges/simplify-the-algebraic-expressions

algebraic.ml

module Lexer = struct
  type token =
      Num of int
    | Id
    | Plus
    | Minus
    | Product
    | Division
    | Exponent
    | LPar
    | RPar

  let is_digit n = n >= '0' && n <= '9'
  and to_digit n = (int_of_char n) - 48

  let scan_number n =
    let rec loop n =
      let c = input_char stdin in
      if is_digit c then
        loop (n * 10 + (to_digit c))
      else (n, c)
    in n |> to_digit |> loop

  let scan peek =
    let rec loop peek = match peek with
      | ' ' | '\t' -> input_char stdin |> loop
      | '\n' -> (None, None)
      | '+' -> (Some Plus, None)
      | '-' -> (Some Minus, None)
      | '*' -> (Some Product, None)
      | '/' -> (Some Division, None)
      | '^' -> (Some Exponent, None)
      | 'x' -> (Some Id, None)
      | '(' -> (Some LPar, None)
      | ')' -> (Some RPar, None)
      | n ->
        if is_digit n then
          let (n, peek) = scan_number n in
          (Some (Num n), Some peek)
        else
          invalid_arg "Unknown token"
    in loop peek

  let scan_all () =
    let rec loop peek =
      let (token, peek) = scan peek in
      match token, peek with
        (Some token, None) -> token :: (loop ' ')
      | (Some token, Some peek) -> token :: (loop peek)
      | (None, _) -> []
    in loop ' '
end

module Parser = struct
  open Lexer
  type element =
      Add of element * element
    | Sub of element * element
    | Prod of element * element
    | Div of element * element
    | Expt of element * element
    | Factor of int * int
  type tree = Node of element * tree list
  let consume input token =
    match input with
      [] -> invalid_arg "Unexpected empty input"
    | x :: tl ->
      if x == token then
        tl
      else invalid_arg "match failed"
  and next input = List.hd input
  and pop input = List.tl input
  and empty input = input = []
  let opt_x input =
    if empty input || next input <> Id then (input, 0)
    else (pop input, 1)
  let factor input =
    match next input with
      Num x ->
        let (input, y) = pop input |> opt_x in
        (input, Factor(x, y))
    | Id -> (pop input, (Factor(1, 1)))
    | _ -> invalid_arg "Unexpected factor"
  let rec e4 input =
    if next input = LPar then
      let input = consume input LPar in
      let (input, e1) = e1 input in
      let input = consume input RPar in
      (input, e1)
    else
      factor input
  and r3 e input =
    if empty input || next input != Exponent then
      (input, e)
    else
      let input = consume input Exponent in
      let (input, e4) = e4 input in
      let (input, r3) = r3 e4 input in
      (input, Expt (e, r3))
  and e3 input =
    let (input, e4) = e4 input in r3 e4 input
  and r2 e input =
    if empty input then (input, e)
    else match next input with
    | Product | Division as token ->
      let input = consume input token in
      let (input, e3) = e3 input in
      let (input, r2) = r2 e3 input in
      if token = Product then
        (input, Prod(e, r2))
      else
        (input, Div(e, r2))
    | _ -> (input, e)
  and e2 input =
    let (input, e3) = e3 input in r2 e3 input
  and r1 e input =
    if empty input then (input, e)
    else match next input with
    | Plus | Minus as token ->
      let input = consume input token in
      let (input, e2) = e2 input in
      let (input, r1) = r1 e2 input in
      if token = Plus then
        (input, Add(e, r1))
      else
        (input, Sub(e, r1))
    | _ -> (input, e)
  and e1 input =
    let (input, e2) = e2 input in r1 e2 input
end

module Evaluator = struct
  open Parser

  let add v1 v2 = Array.map2 (+) v1 v2
  let sub v1 v2 = Array.map2 (-) v1 v2
  let prod v1 v2 =
    let v = Array.make 6 0 in
    for i = 0 to 5 do
      for j = 0 to 5 do
        if i + j <= 5 then
          v.(i + j) <- v.(i + j) + v1.(i) * v2.(j)
      done
    done;
    v
  let div v1 v2 = Array.map (fun x -> x / v2.(0)) v1
  let expt v1 v2 =
    let v = Array.copy v1 in
    begin
      if v.(0) <> 0 then
        v.(0) <- (float_of_int v1.(0)) ** (float_of_int v2.(0)) |> int_of_float
      else
        v.(v2.(0)) <- v1.(1); v.(1) <- 0
    end;
    v

  let rec eval = function
    | Add(e1, e2) ->
      let v1 = eval e1
      and v2 = eval e2 in
      add v1 v2
    | Sub(e1, e2) ->
      let v1 = eval e1
      and v2 = eval e2 in
      sub v1 v2
    | Prod(e1, e2) ->
      let v1 = eval e1
      and v2 = eval e2 in
      prod v1 v2
    | Div(e1, e2) ->
      let v1 = eval e1
      and v2 = eval e2 in
      div v1 v2
    | Expt(e1, e2) ->
      let v1 = eval e1
      and v2 = eval e2 in
      expt v1 v2
    | Factor(n, e) ->
      let v = Array.make 6 0 in
      v.(e) <- n;
      v

  let rec print exp =
    let found = ref false in
    for i = 5 downto 2 do
      let v = exp.(i) in
      if v <> 0 then
        begin
          if !found then
            begin
              begin
                if v > 0 then
                  print_string " + "
                else
                  print_string " - "
              end;
              begin
                if abs v > 1 then
                  Printf.printf "%dx^%d" (abs v) i
                else
                  Printf.printf "x^%d" i
              end;
            end
          else
            begin
              found := true;
              if v = 1 then
                Printf.printf "x^%d" i
              else if v = -1 then
                Printf.printf "-x^%d" i
              else
                Printf.printf "%dx^%d" v i
            end
        end
    done;
    let v = exp.(1) in
    begin
      if v <> 0 then
        begin
          if !found then
            begin
              begin
                if v > 0 then
                  print_string " + "
                else
                  print_string " - "
              end;
              begin
                if abs v > 1 then
                  Printf.printf "%dx" (abs v)
                else
                  print_string "x"
              end
            end
          else
            begin
              found := true;
              if abs v > 1 then
                Printf.printf "%dx" v
              else if v = 1 then
                print_string "x"
              else
                print_string "-x"
            end

        end
    end;
    let v = exp.(0) in
    begin
      if v <> 0 then
        begin
          if !found then
            begin
              begin
                if v > 0 then
                  print_string " + "
                else
                  print_string " - "
              end;
              Printf.printf "%d" (abs v)
            end
          else
            Printf.printf "%d" v
        end
    end;
    print_newline ()
end

let () =
  for i = 1 to read_int () do
    Lexer.scan_all ()
    |> Parser.e1
    |> snd
    |> Evaluator.eval
    |> Evaluator.print
  done