Compute generic-like pool ball triangles

generic_pool_ball_triangles.scala

object GenericPoolBallTriangles {

  def getHeads(values: Seq[Int], m: Int): Set[Seq[Int]] =
    values.combinations(m).map(_.permutations).flatten.toSet

  def fillPuzzle(headNumbers: Seq[Int], m: Int): Seq[Seq[Int]] = {
    def fillPuzzle(values: Seq[Int]): Seq[Seq[Int]] = {
      if(values.size == 1) Seq(values)
      else Seq.concat(
        List(values),
        fillPuzzle(
          (
            for(i <- (0 until values.size - 1)) yield (values(i) - values(i + 1) + m) % m
          ).toList
        )
      )
    }

    fillPuzzle(headNumbers)
  }

  def isOkPuzzle(puzzle: Seq[Seq[Int]], allNumbers: Seq[Int]): Boolean =
    puzzle.flatten.sorted == allNumbers

  // numberOfAllPos: 全ての数字の個数
  // numberOfRepeat: それぞれの数字を何個使って良いか
  def checkMain(numberOfAllPos: Int, numberOfRepeat: Int): Unit = {

    val m = numberOfAllPos / numberOfRepeat
    require(numberOfAllPos % numberOfRepeat == 0)
    val headSize = Stream.from(1).find(j => j * (j + 1) / 2 >= numberOfAllPos).get
    require(headSize * (headSize + 1) / 2 == numberOfAllPos)
    println(headSize, m)
    val values = Seq.fill(numberOfRepeat)(0 until m).flatten.map(j => j % m).sortBy(k => k).toList

    val heads = getHeads(values, headSize)
    val puzzles = heads
      .map(p => fillPuzzle(p, m))
      .filter(p => isOkPuzzle(p, values))
      .toList
    puzzles.foreach(println)
    println(s"count: ${puzzles.size}")
    println("=" * 30)
  }

  def main(args: Array[String]) {
    val args = Seq(
      (1, 1),
      (3, 1),
      (6, 1), (6, 2), (6, 3),
      (10, 1), (10, 2), (10, 5),
      (15, 1), (15, 3), (15, 5),
      // (21, 1),
      (21, 3), (21, 7),
      // (28, 1), (28, 2), (28, 4), (28, 7), (28, 14),
    )
    args.foreach(t => checkMain(t._1, t._2))
  }
}