chess.scala

  /**
    * Rejection rules for a single piece
    * @param ort it blocks vertical and horizontal
    * @param diag it blocks diagonals
    * @param squares it blocks concrete squares: close area for King, two ways to attack for Peasant, 8 squares for Horse
    */
  case class Rules(ort: Boolean, diag: Boolean, squares: (Int, Int)*) {
    /**
      * Cached map of paired rules
      */
    lazy val pair: Map[String, PairRules] = rules.mapValues(r ⇒ PairRules(
      ort = ort || r.ort,
      diag = diag || r.diag,
      edge = this == r,
      squares = squares ++ r.squares.map(xy ⇒ (-xy._1, -xy._2))
    ))
  }

  /**
    * Rejection rules for a pair of placed and to be placed pieces
    * @param ort it blocks orts
    * @param diag it blocks diagonals
    * @param edge whether this piece constructs an edge: not to count pieces twice, we allow the second one to be placed only on the right side of the firts
    * @param squares rejected squares
    */
  case class PairRules(ort: Boolean, diag: Boolean, edge: Boolean, squares: Seq[(Int, Int)])

  lazy val rules: Map[String, Rules] = Map(
    // Peasant can eat one step diag
    "P" → Rules(ort = false, diag = false, squares = (1, 1), (1, -1)),
    // Towers on the orts
    "T" → Rules(ort = true, diag = false),
    // Horses on (1,2) squares
    "H" → Rules(ort = false, diag = false, squares = (for {
      xy ← Seq((1, 2), (2, 1))
      xm ← Seq(1, -1)
      ym ← Seq(1, -1)
    } yield (xy._1 * xm, xy._2 & ym)): _*),
    // Elephant on diagonals
    "E" → Rules(ort = false, diag = true),
    // Queen both orts and diagonals
    "Q" → Rules(ort = true, diag = true),
    // King eats close area nerby it
    "K" → Rules(ort = false, diag = false, squares = (for {
      x ← -1 to 1
      y ← -1 to 1
    } yield (x, y)): _*)
  )

  /**
    * Current situation on the board
    * @param ortX rejected orts on X
    * @param ortY rejected orts on Y
    * @param diagSum rejected sums of X and Y
    * @param diagDiff rejected diffs of X and Y
    * @param squares rejected squares
    * @param edge current edge -- the rightmost one (latest in incoming sequence)
    */
  case class Situation(
                        ortX:     Set[Int]        = Set.empty,
                        ortY:     Set[Int]        = Set.empty,
                        diagSum:  Set[Int]        = Set.empty,
                        diagDiff: Set[Int]        = Set.empty,
                        squares:  Set[(Int, Int)] = Set.empty,
                        edge:     (Int, Int)      = (0, 0)
                      )

  def solve(pieces: Seq[String], width: Int, height: Int): Long = {

    /**
      * Produces a stream of allowed coordinates for the next piece
      * @param add name of the piece to add
      * @param placed already placed pieces
      * @return
      */
    def advance(add: String, placed: Seq[(Rules, Int, Int)] = Seq.empty): Stream[(Rules, Int, Int)] = {
      val situation = placed.foldLeft(Situation()){
        case (s, (r, x, y)) ⇒
          // Getting pair rules from the cache
          val p = r.pair(add)

          Situation(
            ortX = if (r.ort) s.ortX + x else s.ortX,
            ortY = if (r.ort) s.ortY + y else s.ortY,
            diagSum = if (r.diag) s.diagSum + (x + y) else s.diagSum,
            diagDiff = if (r.diag) s.diagDiff + (x - y) else s.diagDiff,
            squares = s.squares ++ p.squares.map(xy ⇒ (xy._1 + x, xy._2 + y)) + ((x, y)),
            edge = if (p.edge) (x, y) else s.edge
          )
      }

      val rr = rules(add)
      for {
        x ← Stream.range(situation.edge._1, width) if !situation.ortX(x) // Remove X orts
        y ← Stream.range(0, height) if !situation.ortY(y) // Remove Y orts
        if !(x == situation.edge._1 && y < situation.edge._2) // For the leftmost edge, reject too small Y
        if !situation.diagSum(x + y) // Not on diag
        if !situation.diagDiff(x - y) // Not on diag
        if !situation.squares((x, y)) // Not in denied squares
      } yield (rr, x, y)
    }

    /**
      * Places all the pieces
      * @param rest rest of the pieces to place
      * @param placed already placed pieces
      * @return number of combinations
      */
    def place(rest: List[String], placed: Seq[(Rules, Int, Int)] = Seq.empty): Long = rest match {
      case Nil ⇒
        1l // Everything is placed
      case add :: Nil ⇒
        advance(add, placed).length.toLong // Last piece
      case add :: tail ⇒
        advance(add, placed).foldLeft(0l){ // For each new situation, calc produce its subsequent situations
          case (s, pos) ⇒ s + place(tail, placed :+ pos)
        }
    }

    place(pieces.toList)
  }

  assert(solve(Seq("P"), 2, 2) == 4, "Any piece on 2x2 gives 4")
  assert(solve(Seq("T"), 3, 3) == 9, "Any piece on 3x3 gives 9")
  assert(solve(Seq("H"), 4, 4) == 16, "Any piece on 4x4 gives 16")
  assert(solve(Seq("E"), 5, 5) == 25, "Any piece on 5x5 gives 25")
  assert(solve(Seq("Q"), 6, 6) == 36, "Any piece on 6x6 gives 36")
  assert(solve(Seq("K"), 7, 7) == 49, "Any piece on 7x7 gives 49")

  assert(solve(Seq("P", "P"), 2, 2) == 4, "4 ways to place two Peasants")
  assert(solve(Seq("T", "T"), 2, 2) == 2, "2 ways to place two Towers")

  assert(solve(Seq("K", "P"), 2, 2) == 0, "No ways to place a King and anything")

  assert(solve(Seq("T", "T", "T"), 3, 3) == 6, "Six ways to place three Towers")

  assert(solve(Seq("E", "E", "E"), 3, 3) == 26, "Twenty six ways to place three Elephants")