Guest User

Premove Checkmate Checker

a guest
Jan 3rd, 2021
313
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #include "testlib.h"
  2. #include <bits/stdc++.h>
  3. using namespace std;
  4. typedef pair<int, int> ii;
  5. typedef long long ll;
  6. typedef long double ld;
  7. #define sz(v) ((int)((v).size()))
  8. #define all(v) (v).begin(),(v).end()
  9.  
  10. namespace Chess {
  11.     const char WHITE_KING = 'K';
  12.     const char WHITE_QUEEN = 'Q';
  13.     const char BLACK_KING = 'X';
  14.  
  15.     const vector<int> dx8 = {-1, -1, -1, 0, 0, 1, 1, 1};
  16.     const vector<int> dy8 = {-1, 0, 1, -1, 1, -1, 0, 1};
  17.  
  18.     struct Cell {
  19.         int x, y;
  20.        
  21.         bool operator == (const Cell& o) const {
  22.             return x == o.x && y == o.y;
  23.         }
  24.  
  25.         bool operator != (const Cell& o) const {
  26.             return x != o.x || y != o.y;
  27.         }
  28.  
  29.         bool operator < (const Cell& o) const {
  30.             return x < o.x || (x == o.x && y < o.y);
  31.         }
  32.     };
  33.  
  34.     struct Move {
  35.         Cell from, to;
  36.     };
  37.  
  38.     struct Board {
  39.         Cell wkPos, wqPos, bkPos;
  40.        
  41.         bool operator == (const Board& o) const {
  42.             return wkPos == o.wkPos && wqPos == o.wqPos && bkPos == o.bkPos;
  43.         }
  44.  
  45.         bool operator < (const Board& o) const {
  46.             return (wkPos < o.wkPos) ||
  47.                    (wkPos == o.wkPos && wqPos < o.wqPos) ||
  48.                    (wkPos == o.wkPos && wqPos == o.wqPos && bkPos < o.bkPos);
  49.         }
  50.     };
  51.  
  52.     struct BoardWithMoveNumber {
  53.         int n;
  54.         Board b;
  55.        
  56.         bool operator == (const BoardWithMoveNumber& o) const {
  57.             return n == o.n && b == o.b;
  58.         }
  59.  
  60.         bool operator < (const BoardWithMoveNumber& o) const {
  61.             return n < o.n || (n == o.n && b < o.b);
  62.         }
  63.     };
  64.  
  65.     bool inRange(const int x, const int y) {
  66.         return 0 <= x && x < 8 && 0 <= y && y < 8;
  67.     }
  68.  
  69.     Cell parseCell(const string& s) {
  70.         assert(sz(s) == 2);
  71.         int x = s[0] - 'a';
  72.         int y = s[1] - '1';
  73.         assert(inRange(x, y));
  74.         return Cell{x, y};
  75.     }
  76.  
  77.     Move parseMove(const string& s) {
  78.         assert(sz(s) == 4);
  79.         Cell from = parseCell(s.substr(0, 2));
  80.         Cell to = parseCell(s.substr(2, 2));
  81.         return Move{from, to};
  82.     }
  83.  
  84.     vector<string> writeBoard(const Board& b) {
  85.         vector<string> s(8, string(8, '.'));
  86.         s[b.wkPos.x][b.wkPos.y] = 'K';
  87.         s[b.wqPos.x][b.wqPos.y] = 'Q';
  88.         s[b.bkPos.x][b.bkPos.y] = 'X';
  89.         return s;
  90.     }
  91.  
  92.     string writeCell(const Cell& c) {
  93.         string res;
  94.         res += char(c.x + 'a');
  95.         res += char(c.y + '1');
  96.         return res;
  97.     }
  98.  
  99.     vector<BoardWithMoveNumber> restoreGame(BoardWithMoveNumber b, const map<BoardWithMoveNumber, BoardWithMoveNumber>& parentBoards) {
  100.         vector<BoardWithMoveNumber> v;
  101.         while (true) {
  102.             v.push_back(b);
  103.             auto it = parentBoards.find(b);
  104.             assert(it != parentBoards.end());
  105.             const BoardWithMoveNumber& newB = it->second;
  106.             if (newB == b) break;
  107.             b = newB;
  108.         }
  109.         reverse(all(v));
  110.         return v;
  111.     }
  112.  
  113.     void writeSingleBoard(const Board& b) {
  114.         cerr << "K" << writeCell(b.wkPos) << " Q" << writeCell(b.wqPos) << " / K" << writeCell(b.bkPos);
  115.     }
  116.  
  117.     string getMove(const Board& b1, const Board& b2) {
  118.         string res;
  119.         int eq = 0;
  120.         if (b1.wkPos == b2.wkPos) eq++; else res = "K" + writeCell(b1.wkPos) + "-" + writeCell(b2.wkPos);
  121.         if (b1.wqPos == b2.wqPos) eq++; else res = "Q" + writeCell(b1.wqPos) + "-" + writeCell(b2.wqPos);
  122.         if (b1.bkPos == b2.bkPos) eq++; else res = "K" + writeCell(b1.bkPos) + "-" + writeCell(b2.bkPos);
  123.         assert(eq == 2);
  124.         assert(!res.empty());
  125.         return res;
  126.     }
  127.  
  128.     vector<string> convertToMoves(const vector<BoardWithMoveNumber>& v) {
  129.         vector<string> moves;
  130.         for (int i = 1; i < sz(v); i++) {
  131.             assert(v[i - 1].n + 1 == v[i].n);
  132.             const string move = getMove(v[i - 1].b, v[i].b);
  133.             moves.push_back(move);
  134.         }
  135.         return moves;
  136.     }
  137.  
  138.     void writeSequence(const BoardWithMoveNumber& last, const map<BoardWithMoveNumber, BoardWithMoveNumber>& parentBoards) {
  139.         vector<BoardWithMoveNumber> v = restoreGame(last, parentBoards);
  140.         vector<string> moves = convertToMoves(v);
  141.         cerr << "Number of moves = " << (sz(moves) + 1) / 2 << endl;
  142.         for (int i = 0, moveId = 1; i < sz(moves); i += 2, moveId++) {
  143.             cerr << moveId << ". ";
  144.             cerr << moves[i];
  145.             if (i + 1 < sz(moves)) {
  146.                 cerr << " " << moves[i + 1];
  147.             }
  148.             cerr << endl;
  149.         }
  150.         cerr << endl;
  151.     }
  152.  
  153.     bool isValidBlackKingMove(const int tox, const int toy, const Board& b) {
  154.         if (!inRange(tox, toy)) return false;
  155.  
  156.         const Cell& wkPos = b.wkPos;
  157.         const Cell& wqPos = b.wqPos;
  158.         const Cell& bkPos = b.bkPos;
  159.  
  160.         // check self position
  161.         if (Cell{tox, toy} == bkPos) return false;
  162.         if (abs(bkPos.x - tox) > 1) return false;
  163.         if (abs(bkPos.y - toy) > 1) return false;
  164.  
  165.         // check white king
  166.         for (int i = -1; i <= 1; i++) {
  167.             for (int j = -1; j <= 1; j++) {
  168.                 if (Cell{wkPos.x + i, wkPos.y + j} == Cell{tox, toy}) return false;
  169.             }
  170.         }
  171.  
  172.         // check white queen
  173.         for (int dir = 0; dir < 8; dir++) {
  174.             for (int len = 1; true; len++) {
  175.                 int qx = wqPos.x + len * dx8[dir];
  176.                 int qy = wqPos.y + len * dy8[dir];
  177.                 if (!inRange(qx, qy)) break;
  178.                 if (Cell{qx, qy} == wkPos) break;
  179.                 if (Cell{qx, qy} == Cell{tox, toy}) return false;
  180.             }
  181.         }
  182.  
  183.         return true;
  184.     }
  185.  
  186.     bool isQueenCaptureMove(const int tox, const int toy, const Board& b) {
  187.         return isValidBlackKingMove(tox, toy, b) && Cell{tox, toy} == b.wqPos;
  188.     }
  189.  
  190.     bool isValidWhiteKingMove(const int tox, const int toy, const Board& b) {
  191.         if (!inRange(tox, toy)) return false;
  192.  
  193.         const Cell& wkPos = b.wkPos;
  194.         const Cell& wqPos = b.wqPos;
  195.         const Cell& bkPos = b.bkPos;
  196.  
  197.         // check self position
  198.         if (Cell{tox, toy} == wkPos) return false;
  199.         if (abs(wkPos.x - tox) > 1) return false;
  200.         if (abs(wkPos.y - toy) > 1) return false;
  201.  
  202.         // check black king
  203.         for (int i = -1; i <= 1; i++) {
  204.             for (int j = -1; j <= 1; j++) {
  205.                 if (Cell{bkPos.x + i, bkPos.y + j} == Cell{tox, toy}) return false;
  206.             }
  207.         }
  208.  
  209.         // check self queen
  210.         if (Cell{tox, toy} == wqPos) return false;
  211.  
  212.         return true;
  213.     }
  214.  
  215.     bool isValidWhiteQueenMove(const int tox, const int toy, const Board& b) {
  216.         if (!inRange(tox, toy)) return false;
  217.  
  218.         const Cell& wkPos = b.wkPos;
  219.         const Cell& wqPos = b.wqPos;
  220.         const Cell& bkPos = b.bkPos;
  221.  
  222.         // check self position
  223.         if (Cell{tox, toy} == wqPos) return false;
  224.  
  225.         // check directions
  226.         bool ok = false;
  227.         for (int dir = 0; dir < 8; dir++) {
  228.             for (int len = 1; true; len++) {
  229.                 int qx = wqPos.x + len * dx8[dir];
  230.                 int qy = wqPos.y + len * dy8[dir];
  231.                 if (!inRange(qx, qy)) break;
  232.                 if (Cell{qx, qy} == wkPos) break; // self king blocks the path
  233.                 if (Cell{qx, qy} == bkPos) assert(false); // can capture black king - can't happen
  234.                 if (Cell{qx, qy} == Cell{tox, toy}) ok = true; // ok, but don't return - check assert in all directions
  235.             }
  236.         }
  237.  
  238.         return ok;
  239.     }
  240.  
  241.     bool blackKingHasNoMoves(const Board& b) {
  242.         const Cell& bkPos = b.bkPos;
  243.  
  244.         for (int dir = 0; dir < 8; dir++) {
  245.             int tox = bkPos.x + dx8[dir];
  246.             int toy = bkPos.y + dy8[dir];
  247.             if (isValidBlackKingMove(tox, toy, b)) return false;
  248.         }
  249.         return true;
  250.     }
  251.  
  252.     bool blackKingUnderCheck(const Board& b) {
  253.         const Cell& wkPos = b.wkPos;
  254.         const Cell& wqPos = b.wqPos;
  255.         const Cell& bkPos = b.bkPos;
  256.  
  257.         for (int dir = 0; dir < 8; dir++) {
  258.             for (int len = 1; true; len++) {
  259.                 int qx = wqPos.x + len * dx8[dir];
  260.                 int qy = wqPos.y + len * dy8[dir];
  261.                 if (!inRange(qx, qy)) break;
  262.                 if (Cell{qx, qy} == wkPos) break;
  263.                 if (Cell{qx, qy} == bkPos) return true;
  264.             }
  265.         }
  266.         return false;
  267.     }
  268.  
  269.     bool isCheckmate(const Board& b) {
  270.         return blackKingHasNoMoves(b) && blackKingUnderCheck(b);
  271.     }
  272.  
  273.     bool isStalemate(const Board& b) {
  274.         return blackKingHasNoMoves(b) && !blackKingUnderCheck(b);
  275.     }
  276.  
  277.     bool isValidWhiteMove(const int fromx, const int fromy, const int tox, const int toy, const Board& b) {
  278.         if (Cell{fromx, fromy} == b.wkPos) return isValidWhiteKingMove(tox, toy, b);
  279.         if (Cell{fromx, fromy} == b.wqPos) return isValidWhiteQueenMove(tox, toy, b);
  280.         return false;
  281.     }
  282.  
  283.     void makeMove(const int fromx, const int fromy, const int tox, const int toy, Board& b) {
  284.         if (Cell{fromx, fromy} == b.wkPos) {
  285.             assert(isValidWhiteKingMove(tox, toy, b));
  286.             b.wkPos = Cell{tox, toy};
  287.             return;
  288.         }
  289.         if (Cell{fromx, fromy} == b.wqPos) {
  290.             assert(isValidWhiteQueenMove(tox, toy, b));
  291.             b.wqPos = Cell{tox, toy};
  292.             return;
  293.         }
  294.         if (Cell{fromx, fromy} == b.bkPos) {
  295.             assert(isValidBlackKingMove(tox, toy, b));
  296.             assert(!isQueenCaptureMove(tox, toy, b));
  297.             b.bkPos = Cell{tox, toy};
  298.             return;
  299.         }
  300.         assert(false);
  301.     }
  302.  
  303.     // (error, max number of moves)
  304.     pair<string, int> checkSequence(const vector<string>& inputSequence, const string& wkStartPosStr, const string& wqStartPosStr) {
  305.         //cerr << "sz(inputSequence) == " << sz(inputSequence) << endl;
  306.  
  307.         vector<Move> moves;
  308.         for (const string& moveStr : inputSequence) {
  309.             moves.push_back(parseMove(moveStr));
  310.         }
  311.        
  312.         map<BoardWithMoveNumber, BoardWithMoveNumber> parentBoards;
  313.        
  314.         const Cell wkStartPos = parseCell(wkStartPosStr);
  315.         const Cell wqStartPos = parseCell(wqStartPosStr);
  316.         assert(wkStartPos != wqStartPos);
  317.        
  318.         set<BoardWithMoveNumber> currentBoards;
  319.         // top-right corner
  320.         for (int bkx = wqStartPos.x + 1; bkx <= 7; bkx++) {
  321.             for (int bky = wqStartPos.y + 1; bky <= 7; bky++) {
  322.                 Cell bkStartPos{bkx, bky};
  323.                 if (max(abs(bkStartPos.x - wkStartPos.x), abs(bkStartPos.y - wkStartPos.y)) <= 1) {
  324.                     // kings are too close
  325.                     continue;
  326.                 }
  327.                 Board b;
  328.                 b.wkPos = wkStartPos;
  329.                 b.wqPos = wqStartPos;
  330.                 b.bkPos = bkStartPos;
  331.                 if (blackKingUnderCheck(b)) {
  332.                     continue;
  333.                 }
  334.                 //cerr << "Will add start black position: " << writeCell(bkStartPos) << endl;
  335.                 currentBoards.insert({1, b});
  336.                 parentBoards[{1, b}] = {1, b};
  337.             }
  338.         }
  339.         assert(!currentBoards.empty());
  340.        
  341.         set<BoardWithMoveNumber> checkmates;
  342.         for (const Move& move : moves) {
  343.             /*cerr << "Before move " << writeCell(move.from) << "-" << writeCell(move.to) << endl;
  344.             cerr << "Number of boards = " << sz(currentBoards) << endl;
  345.             for (const BoardWithMoveNumber& bb : currentBoards) {
  346.                 writeSingleBoard(bb.b);
  347.                 cerr << endl;
  348.             }
  349.             cerr << "=======================" << endl;*/
  350.             set<BoardWithMoveNumber> newBoards;
  351.             for (const BoardWithMoveNumber& bb : currentBoards) {
  352.                 const Board& b = bb.b;
  353.                 const int n = bb.n;
  354.                 if (isValidWhiteMove(move.from.x, move.from.y, move.to.x, move.to.y, b)) {
  355.                     Board b2 = b;
  356.                     makeMove(move.from.x, move.from.y, move.to.x, move.to.y, b2);
  357.                     parentBoards[{n + 1, b2}] = {n, b};
  358.                     if (isCheckmate(b2)) {
  359.                         checkmates.insert({n + 1, b2});
  360.                         continue;
  361.                     }
  362.                     if (isStalemate(b2)) {
  363.                         writeSequence({n + 1, b2}, parentBoards);
  364.                         //cerr << "Stalemate" << endl;
  365.                         return {"Stalemate", -1};
  366.                     }
  367.                     const Cell& bkPos = b2.bkPos;
  368.                     for (int dir = 0; dir < 8; dir++) {
  369.                         int tox = bkPos.x + dx8[dir];
  370.                         int toy = bkPos.y + dy8[dir];
  371.                         if (isQueenCaptureMove(tox, toy, b2)) {
  372.                             writeSequence({n + 1, b2}, parentBoards);
  373.                             //cerr << "Queen is captured" << endl;
  374.                             return {"Queen is captured", -1};
  375.                         }
  376.                         if (isValidBlackKingMove(tox, toy, b2)) {
  377.                             Board b3 = b2;
  378.                             makeMove(bkPos.x, bkPos.y, tox, toy, b3);
  379.                             parentBoards[{n + 2, b3}] = {n + 1, b2};
  380.                             newBoards.insert({n + 2, b3});
  381.                         }
  382.                     }
  383.                 } else {
  384.                     newBoards.insert(bb);
  385.                 }
  386.             }
  387.             currentBoards = newBoards;
  388.         }
  389.        
  390.         if (checkmates.empty()) {
  391.             //cerr << "Couldn't give even a single checkmate" << endl;
  392.             return {"Couldn't give even a single checkmate", -1};
  393.         }
  394.         if (!currentBoards.empty()) {
  395.             //cerr << "Couldn't give checkmate in some situations" << endl;
  396.             return {"Couldn't give checkmate in some situations", -1};
  397.         }
  398.        
  399.         vector<int> numbersOfMoves;
  400.         for (const BoardWithMoveNumber& bb : checkmates) {
  401.             numbersOfMoves.push_back(bb.n / 2);
  402.         }
  403.         const int maxNumberOfMoves = *max_element(all(numbersOfMoves));
  404.         if (maxNumberOfMoves > 50) {
  405.             return {"Max number of moves is greater than 50", maxNumberOfMoves};
  406.         }
  407.  
  408.         /*cerr << "Success!" << endl;
  409.         cerr << "Numbers of moves: ";
  410.         for (const int x : numbersOfMoves) cerr << x << " ";
  411.         cerr << endl;
  412.        
  413.         cerr << "Gave " << sz(checkmates) << " checkmates" << endl;
  414.         for (const BoardWithMoveNumber& bb : checkmates) {
  415.             writeSequence(bb, parentBoards);
  416.         }*/
  417.         return {"", maxNumberOfMoves}; // ok
  418.     }
  419. }
  420.  
  421. vector<string> readSequence(InStream& stream, TResult waResult) {
  422.     vector<string> inputSequence;
  423.     while (!stream.seekEof()) {
  424.         const string token = stream.readToken();
  425.         if (sz(token) != 4) stream.quitf(waResult, "Wrong move string length: %d", sz(token));
  426.         if (token[0] < 'a' || token[0] > 'h') stream.quitf(waResult, "Wrong 0-th character: '%c'", token[0]);
  427.         if (token[1] < '1' || token[1] > '8') stream.quitf(waResult, "Wrong 1-st character: '%c'", token[1]);        
  428.         if (token[2] < 'a' || token[2] > 'h') stream.quitf(waResult, "Wrong 2-nd character: '%c'", token[2]);
  429.         if (token[3] < '1' || token[3] > '8') stream.quitf(waResult, "Wrong 3-rd character: '%c'", token[3]);        
  430.         if (token[0] == token[2] && token[1] == token[3]) stream.quitf(waResult, "Wrong move: start and finish cells are equal");
  431.         inputSequence.push_back(token);
  432.         if (sz(inputSequence) > 500) {
  433.             stream.quitf(waResult, "Too many moves in the sequence");
  434.         }
  435.     }
  436.     if (inputSequence.empty()) {
  437.         stream.quitf(waResult, "Sequence is empty");
  438.     }
  439.     return inputSequence;
  440. }
  441.  
  442. vector<string> getAlexInputSequence() {
  443.     return {
  444.         // Kc3, Qd4
  445.         "c3d3", "d3e3", "e3f3", "f3g3",
  446.         // Kg3, Qd4
  447.         "g3g4", "g3h4", "h4g4", "d4c5", "g4g3", "h4g3",
  448.         // Kg3, Qc5
  449.         "g3f3", "f3e3", "e3d3", "d3c4", "c4b5", "b5b6", "b6b7",
  450.         // Kb7, Qc5
  451.         "b7c6", "b7b6", "b6c6",
  452.         // Kc6, Qc5, black is either e+ or {c8 d8}
  453.         "c5d5", "c6c7", "d5a5", "a5c7", "c7h7", "h7b7",
  454.         // Kc7, Qa5 (or already mated)
  455.         "c7c6", "a5d5",
  456.         // Kc6, Qd5
  457.         "c6c5", "c5d4", "d4e4", "e4f4", "f4g4",
  458.         // Kg4, Qd5
  459.         "g4g5", "g4h5", "h5g5", "d5c6", "g5g4", "h5g4",
  460.         // Kg4, Qc6
  461.         "g4f4", "f4e4", "e4d4", "d4c5", "c5b6", "b6b7",
  462.         // Kb7, Qc6
  463.         "b7c7", "b7b8", "b8c7", "c6d6",
  464.         // Kc7, Qd6
  465.         "c7d7", "c7c8", "c8d7", "d6e6",
  466.         // Kd7, Qe6
  467.         "d7e7", "d7d8", "d8e7",
  468.         // Ke7, Qe6
  469.         "e6g4", "e7f7", "g4h4",
  470.     };
  471. }
  472.  
  473. vector<string> getSergeyInputSequence() {
  474.     return {
  475.         // Move king to h4
  476.         "c3d3", "d3e3", "e3f3", "f3g3", "g3h4",
  477.         // If black king is on g5 or h5, place queen to h4, then to e7 and give checkmate
  478.         "d4h4", "h4e7", "g3g4", "g3f3", "f3g4", "g4f5", "g4f4", "f4f5", "e7g5", "f5f6", "g5g7",
  479.         // Our king is on h4, go to f5, then make triangle f6-g6-f5 and try to place queen to f6
  480.         "h4g4", "h4h5", "h5g4", "g4f5", "g4f4", "f4f5", "f5f6", "f6g6", "g6f5", "d4f6",
  481.         // If we succeeded in placing queen is on f6, white king is on f5, and black king is in e7..h8 area
  482.         // Move queen to b6, king to b7 and give easy checkmate
  483.         "f6b6", "f5e5", "e5d5", "d5c5", "c5b5", "b5a6", "a6b7", "a6a7", "a7b7",
  484.         "b6c6", "b7c7", "b7b8", "b8c7",
  485.         "c6d6", "c7d7", "c7c8", "c8d7",
  486.         "d6e6", "d7e7", "d7d8", "d8e7",
  487.         "e6g4", "e7f7", "g4h4",
  488.         // If we couldn't place queen on f6, white king is on f6, and black king is on h5, h6, h7
  489.         // Place queen to g file and give checkmate
  490.         "d4g1", "g1g2", "f6f5", "g2g4", "g4g5", "f5f6", "g5g7"
  491.     };
  492. }
  493.  
  494. int main(int argc, char* argv[]) {
  495.     /*
  496.     const vector<string> inputSequence = getSergeyInputSequence();
  497.     cerr << "sz = " << sz(inputSequence) << endl;
  498.     const string error = Chess::checkSequence(inputSequence, "c3", "d4");
  499.     cerr << "Error: " << error << endl;
  500.     */
  501.  
  502.     registerTestlibCmd(argc, argv);
  503.  
  504.     const string wkStartPosStr = inf.readToken();
  505.     const string wqStartPosStr = inf.readToken();
  506.     assert(
  507.         ((wkStartPosStr == "c3") && (wqStartPosStr == "d4"))
  508.         ||
  509.         ((wkStartPosStr == "f6") && (wqStartPosStr == "e5"))
  510.     );
  511.  
  512.     const vector<string> jInputSequence = readSequence(ans, _fail);
  513.     const pair<string, int> jError = Chess::checkSequence(jInputSequence, wkStartPosStr, wqStartPosStr);
  514.     if (!jError.first.empty()) {
  515.         quitf(_fail, "Jury solution is wrong: %s", jError.first.c_str());
  516.     }
  517.     assert(jError.second > 0);
  518.  
  519.     const vector<string> pInputSequence = readSequence(ouf, _wa);
  520.     const pair<string, int> pError = Chess::checkSequence(pInputSequence, wkStartPosStr, wqStartPosStr);
  521.     if (!pError.first.empty()) {
  522.         quitf(_wa, "Participant solution is wrong: %s", pError.first.c_str());
  523.     }
  524.     assert(pError.second > 0);
  525.  
  526.     quitf(_ok, "Correct solution, length %d, max number of moves %d", sz(pInputSequence), pError.second);
  527. }
  528.  
RAW Paste Data