Home Posts Tags Post Search Tag Search

Post 48

Hacker Rank 22 - Tree Functions

Published on: 2025-07-17 Tags: elixir, Side Project, Hacker Rank
defmodule Solution do
  def run(print \\ true) do
    read()
    |> print(print, "read")
    |> event_helper()
  end

  # Reads input from stdin, splits it into steps
  def read do
    IO.read(:stdio, :all)
    |> String.split("\n", trim: true)
    |> Enum.map(&String.split(&1, " ", trim: true))
  end

  # Initializes the tree and processes each event
  def event_helper([_ | steps]) do
    Enum.reduce(
      steps,
      %{
        tree: %{
          0 => %{
            value: 0,
            left: nil,
            right: nil,
            parent: nil,
            children: []
          }
        },
        current: 0,
        next_id: 1
      },
      fn step, acc ->
        event(step, acc)
      end
    )
  end

  # Change the value of the current node
  def event(["change", val], %{current: node} = acc) do
    value = String.to_integer(val)
    put_in(acc, [:tree, node, :value], value)
  end

  # Print the value of the current node
  def event(["print"], %{current: node} = acc) do
    IO.puts(acc.tree[node].value)
    acc
  end

  # Move to the left sibling of the current node
  def event(["visit", "left"], %{current: node, tree: tree} = acc) do
    put_in(acc, [:current], tree[node].left)
  end

  # Move to the right sibling of the current node
  def event(["visit", "right"], %{current: node, tree: tree} = acc) do
    put_in(acc, [:current], tree[node].right)
  end

  # Move to the parent of the current node
  def event(["visit", "parent"], %{current: node, tree: tree} = acc) do
    put_in(acc, [:current], tree[node].parent)
  end

  # Visit the Nth child (1-based index) of the current node
  def event(["visit", "child", index], %{current: node, tree: tree} = acc) do
    idx = String.to_integer(index) - 1
    new_id = Enum.at(tree[node].children, idx)
    put_in(acc, [:current], new_id)
  end

  # Insert a new node to the left of the current node
  def event(["insert", "left", val], %{current: node, next_id: id, tree: tree} = acc) do
    value = String.to_integer(val)
    parent = tree[node].parent

    new_node = %{
      value: value,
      parent: parent,
      children: [],
      left: tree[node].left,
      right: node
    }

    acc =
      if new_node.left != nil do
        put_in(acc, [:tree, new_node.left, :right], id)
      else
        acc
      end

    acc = put_in(acc, [:tree, node, :left], id)

    index = Enum.find_index(tree[parent].children, &(&1 == node))
    children = List.insert_at(tree[parent].children, index, id)

    acc = put_in(acc, [:tree, parent, :children], children)
    put_in(acc, [:tree, id], new_node) |> Map.update!(:next_id, &(&1 + 1))
  end

  # Insert a new node to the right of the current node
  def event(["insert", "right", val], %{current: node, next_id: id, tree: tree} = acc) do
    value = String.to_integer(val)
    parent = tree[node].parent

    new_node = %{
      value: value,
      parent: parent,
      children: [],
      left: node,
      right: tree[node].right
    }

    acc =
      if new_node.right != nil do
        put_in(acc, [:tree, new_node.right, :left], id)
      else
        acc
      end

    acc = put_in(acc, [:tree, node, :right], id)

    index = Enum.find_index(tree[parent].children, &(&1 == node))
    children = List.insert_at(tree[parent].children, index + 1, id)

    acc = put_in(acc, [:tree, parent, :children], children)
    put_in(acc, [:tree, id], new_node) |> Map.update!(:next_id, &(&1 + 1))
  end

  # Insert a new child node as the first child of the current node
  def event(["insert", "child", val], %{current: node, next_id: id, tree: tree} = acc) do
    value = String.to_integer(val)
    old_first = List.first(tree[node].children)

    new_node = %{
      value: value,
      parent: node,
      children: [],
      left: nil,
      right: old_first
    }

    acc =
      if old_first do
        put_in(acc, [:tree, old_first, :left], id)
      else
        acc
      end

    children = [id | tree[node].children]
    acc = put_in(acc, [:tree, node, :children], children)
    put_in(acc, [:tree, id], new_node) |> Map.update!(:next_id, &(&1 + 1))
  end

  # Delete the current node and all its descendants
  def event(["delete"], %{current: node, tree: tree} = acc) do
    parent = tree[node].parent
    left = tree[node].left
    right = tree[node].right

    # Recursive function to collect all descendants
    collect = fn collect, id ->
      Enum.reduce(tree[id].children, [id], fn c, acc ->
        acc ++ collect.(collect, c)
      end)
    end

    to_delete = collect.(collect, node)
    new_tree = Map.drop(tree, to_delete)

    # Remove node from parent's children list
    updated_children = List.delete(tree[parent].children, node)
    new_tree = put_in(new_tree, [parent, :children], updated_children)

    # Update left and right sibling pointers
    new_tree =
      if left, do: put_in(new_tree, [left, :right], right), else: new_tree
    new_tree =
      if right, do: put_in(new_tree, [right, :left], left), else: new_tree

    %{acc | tree: new_tree, current: parent}
  end

  # Optional print helper for debugging
  defp print(data, true, label), do: IO.inspect(data, label: label)
  defp print(data, false, _), do: data
end

Solution.run(false)