Scala program for my CS 113 class that generates LaTeX code for the Heapify problems in HW #14

hw14.scala

import cats._, cats.data._, cats.syntax.all._
import scala.annotation.tailrec

object Main {
  
  case class Snapshot(comparisons: Int, exchanges: Int, tikz: String)

  def main(args: Array[String]): Unit = {
    val vector: Vector[Int] = Vector(55, 50, 10, 40, 80, 90, 60, 100, 70, 80, 20, 50, 22)
    val (snapshots, _, ()) = heapify[Int].runEmpty(vector).value
    val result = snapshots.foldLeft(("", 0, 0)) {
      case ((str, cumComp, cumExch), Snapshot(comparisons, exchanges, tikz)) =>
        val newCumComp = cumComp + comparisons
        val newCumExch = cumExch + exchanges
        val caption = s"$comparisons comparisons, $newCumComp cum. comparisons, $exchanges exchanges, $newCumExch cum. exchanges."
        val newText = s"\\begin{figure}[!ht]\n\\centering\n$tikz\n\\caption{$caption}\n\\end{figure}\n"
        (str ++ newText, newCumComp, newCumExch)
    }
    println(result._1)
  }

  def heapify[A : PartialOrder : Show]: ReaderWriterState[Vector[A], Vector[Snapshot], Vector[A], Unit] =
    ReaderWriterState.ask[Vector[A], Vector[Snapshot], Vector[A]].flatMap { input =>
      (0 until input.size).toVector.traverse_(insertIntoHeap[A])
    }

  def insertIntoHeap[A : PartialOrder : Show](envIndex: Int): ReaderWriterState[Vector[A], Vector[Snapshot], Vector[A], Unit] =
    ReaderWriterState { (input, heap) =>
      val envElem = input(envIndex)

      val firstIndex = heap.size
      val toAdd = heap :+ envElem

      @tailrec def go(heap: Vector[A], heapifyIndex: Int, comparisons: Int = 0, exchanges: Int = 0): (Vector[A], Snapshot) = {
        val parentIndex = (heapifyIndex - 1) / 2

        val hElem = heap(heapifyIndex)
        val pElem = heap(parentIndex)
        if(hElem > pElem) {
          go(heap.updated(heapifyIndex, pElem).updated(parentIndex, hElem), parentIndex, comparisons + 1, exchanges + 1)
        } else {
          (heap, Snapshot(comparisons + 1, exchanges, heapTikz[A](heap)))
        }
      }

      val (result, snapshot) = go(toAdd, firstIndex)
      (Vector(snapshot), result, ())
    }
  
  def heapTikz[A : Show](heap: Vector[A]): String = {
    val start = "\\begin{tikzpicture}[heap] \\"
    val end = ";\\end{tikzpicture}"
    def getNode(index: Int): String = {
      val myValue: String = heap(index).show
      val myLeft = (index * 2) + 1
      val myRight = (index * 2) + 2
      val myLeftVal = Option.when(myLeft < heap.size)(getNode(myLeft))
      val myRightVal = Option.when(myRight < heap.size)(getNode(myRight))
      val myLeftString = myLeftVal.fold("")(s => s"child{$s}")
      val myRightString = myRightVal.fold("")(s => s"child{$s}")
      s"node{$myValue} $myLeftString $myRightString"
    }
    start ++ getNode(0) ++ end
  }

}