Generates formulae for sum of powers.

sum-of-power.scala

#!/bin/bash

exec /opt/scala-2.10/bin/scala -feature "$0" "$@"
!#
import scala.language.postfixOps

/**
 * From https://stackoverflow.com/questions/25129721/scala-memoization-how-does-this-scala-memo-work.
 *
 * Generic way to create memoized functions (even recursive and multiple-arg ones)
 *
 * @param f the function to memoize
 * @tparam I input to f
 * @tparam K the keys we should use in cache instead of I
 * @tparam O output of f
 */
case class Memo[I <% K, K, O](f: I => O) extends (I => O) {
  import scala.collection.mutable.{Map => Dict}

  type Input = I
  type Key = K
  type Output = O

  val cache = Dict.empty[K, O]

  override def apply(x: I) = cache getOrElseUpdate(x, f(x))
}

object Memo {
  /**
   * Type of a simple memoized function e.g. when I = K
   */
  type ==>[I, O] = Memo[I, I, O]
}

implicit class _RichInt(n: Int) {
  private def fact(n: Long): BigInt = {
    lazy val f: Memo.==>[(BigInt, Long), BigInt] = Memo {
      case (a, 0) => a
      case (a, 1) => a
      case (a, n) => f(a * n, n - 1)
    }

    f(BigInt(1), n)
  }

  def ! = fact(n)
  def P(k: Long): BigInt = fact(n) / fact(n - k)
  def C(k: Long): Long = (P(k) / fact(k)).toLong
}

def gcd(lhs: Long, rhs: Long): Long = {
  if (rhs == 0) {
    lhs.abs
  } else {
    gcd(rhs, lhs % rhs)
  }
}

def lcm(lhs: Long, rhs: Long): Long = {
  (lhs * rhs).abs / gcd(lhs, rhs)
}

object Primes {
  private lazy val notDivisibleBy2: Stream[Long] = 3L #:: notDivisibleBy2.map(_ + 2)
  private lazy val notDivisibleBy2Or3: Stream[Long] = notDivisibleBy2
      .grouped(3)
      .map(_.slice(1, 3))
      .flatten
      .toStream
  private lazy val notDivisibleBy2Or3Or5: Stream[Long] = notDivisibleBy2Or3
      .grouped(10)
      .map { g =>
        g.slice(1, 7) ++ g.slice(8, 10)
      }
      .flatten
      .toStream

  lazy val primes: Stream[Long] = 2L #::
      notDivisibleBy2.head #::
      notDivisibleBy2Or3.head #::
      notDivisibleBy2Or3Or5.filter { i =>
        i < 49 || primes.takeWhile(_ <= Math.sqrt(i).toLong).forall(i % _ != 0)
      }

  def <(n: Long): Stream[Long] = primes.takeWhile(_ <= n)

  def largestPower(n: Long, p: Long): Int = {
    @tailrec
    def _largestPower(x: Int, n: Long): Int = {
      if (n % p != 0) {
        x
      } else {
        _largestPower(x + 1, n / p)
      }
    }

    _largestPower(0, n)
  }

  def factors(n: Long): Map[Long, Int] = {
    val primeFactorsOfN = (Primes < n).filter(n % _ == 0)

    primeFactorsOfN.map { p =>
      (p, largestPower(n, p))
    }.toMap
  }
}

case class Polynomial(val coefficients: Long*) {
  // TODO: Factor the polynomial.
  override def toString: String = {
    def toTermString(coefficient: Long, exponent: Int): String = {
      val coefficientPart = if (coefficient == 1) {
        ""
      } else {
        coefficient.toString
      }
      val variablePart = if (exponent == 0) {
        ""
      } else if (exponent == 1) {
        "n"
      } else {
        s"n^${exponent}"
      }

      s"${coefficientPart}${variablePart}"
    }

    val coefficientExponentPairs = coefficients
        .zipWithIndex
        .reverse
        .filter(_._1 != 0)
    val head = coefficientExponentPairs.head
    val tail = coefficientExponentPairs.tail

    val firstTerm = toTermString(head._1, head._2)
    tail
        .foldLeft(firstTerm) { case (lhs, rhs) =>
          val (coefficient, exponent) = rhs
          val operator = if (coefficient < 0) {
            "-"
          } else {
            "+"
          }

          s"${lhs} ${operator} ${toTermString(coefficient.abs, exponent)}"
        }
  }

  def apply(n: Int): BigInt = {
    coefficients
        .zipWithIndex
        .map { case (coefficient, i) =>
          coefficient * BigInt(n).pow(i)
        }.sum
  }

  def -(that: Polynomial): Polynomial = {
    val length = Math.max(coefficients.length, that.coefficients.length)
    val lhs = coefficients.padTo(length, 0L)
    val rhs = that.coefficients.padTo(length, 0L)

    Polynomial(lhs.zip(rhs).map { case (l, r) =>
      l - r
    } :_*)
  }

  def /(n: Long): Polynomial = {
    Polynomial(coefficients.map(_ / n) :_*)
  }
}

implicit class _RichLongWithPolynomial(n: Long) {
  def *(polynomial: Polynomial): Polynomial = {
    Polynomial(polynomial.coefficients.map(n * _) :_*)
  }
}

class Σp(val numerator: Polynomial, val denominator: Long) {
  override def toString: String = {
    s"(${numerator}) / ${denominator}"
  }

  def -(rhs: Σp): Σp = {
    val d = lcm(denominator, rhs.denominator)

    Σp(
      d / denominator * numerator - d / rhs.denominator * rhs.numerator,
      d)
  }

  def /(rhs: Long): Σp = {
    Σp(numerator, rhs * denominator)
  }

  def apply(n: Int): BigInt = {
    numerator(n) / denominator
  }
}

object Σp {
  def apply(power: Int): Σp = {
    lazy val f: Memo.==>[Int, Σp] = Memo {
      case 0 => Σp(Polynomial(0L, 1L), 1L)
      case p => {
        val init = new Σp(
          Polynomial(0L +: (1 to p + 1).map { i =>
            p + 1 C i
          } :_*),
          1)

        (0 to p - 1).map { i =>
          (p + 1 C i) * f(i)
        }.foldLeft(init) { (lhs, rhs) =>
          lhs - rhs
        } / (p + 1)
      }
    }

    f(power)
  }

  def apply(numerator: Polynomial, denominator: Long): Σp = {
    val g = numerator.coefficients.foldLeft(denominator) { (lhs, rhs) =>
      gcd(lhs, rhs)
    }

    new Σp(numerator / g, denominator / g)
  }
}

implicit class _RichLongWithΣp(n: Long) {
  def *(σp: Σp): Σp = {
    val g = gcd(n, σp.denominator)

    Σp(n / g * σp.numerator, σp.denominator / g)
  }
}

if (args.length == 0) {
  (0 to 48)
      .foreach { p =>
        println(Σp(p))
      }
} else {
  args.map { arg =>
    val σp = Σp(arg.toInt)

    (0 to 3)
        .map(σp(_))
        .foldLeft(σp.toString) { (lhs, rhs) =>
          s"${lhs}\n${rhs}"
        }
  }.foreach(println)
}