(C++) KenKen Puzzle Solver

kenken.cpp

/*
  Program that solves a kenken puzzle and pretty prints the result.
  Reading the puzzle as input is not currently supported.
*/

#include <fcntl.h>
#include <locale.h>
#include <iomanip>
#include <iostream>
#include <locale>
#include <sstream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>

typedef std::vector<std::vector<int>> Board;

namespace box {
const wchar_t N = L' ';        // None
const wchar_t V = L'\u2503';   // Vertical
const wchar_t H = L'\u2501';   // Horizontal
const wchar_t VH = L'\u254B';  // Cross
const wchar_t UR = L'\u2517';  // Bottom left corner
const wchar_t DR = L'\u250F';  // Top left corner
const wchar_t UL = L'\u251B';  // Bottom right corner
const wchar_t DL = L'\u2513';  // Top right corner
const wchar_t HD = L'\u2533';  // Horizontal / Down
const wchar_t HU = L'\u253B';  // Horizontal / Up
const wchar_t VR = L'\u2523';  // Vertical / Right
const wchar_t VL = L'\u252B';  // Vertical / Left
}  // namespace box

class Position {
   public:
    int x;
    int y;
    Position(int _x, int _y) : x(_x), y(_y) {}

    inline bool operator==(const Position& rhs) const {
        return this->x == rhs.x && this->y == rhs.y;
    }

    bool sharesNorthEdge(const Position& p) const {
        return this->x == p.x && this->y == p.y + 1;
    }
    bool sharesSouthEdge(const Position& p) const {
        return this->x == p.x && this->y == p.y - 1;
    }
    bool sharesEastEdge(const Position& p) const {
        return this->y == p.y && this->x == p.x - 1;
    }
    bool sharesWestEdge(const Position& p) const {
        return this->y == p.y && this->x == p.x + 1;
    }
};

namespace std {
template <>
struct hash<Position> {
    std::size_t operator()(const Position& p) const {
        return (p.x << 16) ^ p.y;
    }
};
}  // namespace std

class Cage {
    std::vector<Position> positions;
    int value;
    wchar_t symbol;
    Position topLeft;

   public:
    Cage(std::vector<Position> _positions, int _value, wchar_t _symbol)
        : positions(_positions),
          value(_value),
          symbol(_symbol),
          topLeft(-1, -1) {
        // Validate constraint definitions
        if (!(_symbol == '*' || _symbol == '/' || _symbol == '+' ||
              _symbol == '-' || _symbol == ' ')) {
            throw std::invalid_argument("Invalid operation symbol passed.");
        }
        if (_symbol == ' ' && this->positions.size() != 1) {
            throw std::invalid_argument(
                "Illegal cage size for that operation.");
        }
        if ((_symbol == '/' || _symbol == '-') && this->positions.size() != 2) {
            throw std::invalid_argument(
                "Illegal cage size for that operation.");
        }

        for (const Position& p : this->positions) {
            if (topLeft.x == -1 || topLeft.y == -1) {
                topLeft.x = p.x;
                topLeft.y = p.y;
                continue;
            }
            if (p.x < topLeft.x || (p.x == topLeft.x && p.y < topLeft.y)) {
                topLeft.x = p.x;
                topLeft.y = p.y;
            }
        }
    }

    const Position& getRenderPosition() const { return this->topLeft; }
    // Returns true if this cage contains this position.
    bool hasPosition(int x, int y) const {
        Position q = Position(x, y);
        for (const Position& p : this->positions) {
            if (p == q) return true;
        }
        return false;
    }

    // Returns whether or not a value fits this cage, given
    // the current board.
    bool fits(const Board& board, int input) const {
        std::vector<int> values;
        int emptySquares = 0;
        for (const Position& p : this->positions) {
            int v = board[p.x][p.y];
            if (!v)
                ++emptySquares;
            else
                values.push_back(v);
        }
        if (emptySquares > 1) {
            return true;
        }
        int acc;
        switch (this->symbol) {
            case ' ':
                return (input == this->value);
            case '*':
                acc = 1;
                for (int v : values) acc *= v;
                return (acc * input) == this->value;
            case '-':
                return ((input - values[0]) == this->value) ||
                       ((values[0] - input) == this->value);
            case '+':
                acc = 0;
                for (int v : values) acc += v;
                return (acc + input) == this->value;
            case '/':
                return ((input / values[0]) == this->value &&
                        (input % values[0] == 0)) ||
                       ((values[0] / input) == this->value &&
                        (values[0] % input == 0));
        }

        return false;
    }

    const std::vector<Position>& getPositions() const {
        return this->positions;
    }
    friend std::wostream& operator<<(std::wostream&, Cage const&);
};

std::wostream& operator<<(std::wostream& os, Cage const& c) {
    return os << std::to_wstring(c.value) + c.symbol;
}

struct EdgeSet {
    bool north;
    bool east;
    bool south;
    bool west;

    EdgeSet() : north(false), east(false), south(false), west(false) {}

    EdgeSet(bool v) : north(v), east(v), south(v), west(v) {}

    EdgeSet(bool n, bool e, bool s, bool w)
        : north(n), east(e), south(s), west(w) {}

    wchar_t toCorner() const {
        using namespace box;
        if (north && east && south && west) return N;
        if (!north && !east && !south && !west) return VH;

        // edge cases
        if (!north && east && !south && west) return V;
        if (north && !east && south && !west) return H;

        // 3 prong cases
        if (!north && east && !south && !west) return VR;
        if (!north && !east && !south && west) return VL;
        if (north && !east && !south && !west) return HD;
        if (!north && !east && south && !west) return HU;

        // corner cases
        if (!north && !east && south && west) return UL;
        if (north && !east && !south && west) return DL;
        if (!north && east && south && !west) return UR;
        if (north && east && !south && !west) return DR;

        return L'?';
    }
};

class Puzzle {
   private:
    std::vector<Cage> cages;
    Board board;
    std::unordered_map<Position, Cage*> cageMap;
    std::unordered_map<Position, EdgeSet> innerEdges;

    bool isInitialized() const {
        for (int i = 0; i < board.size(); ++i) {
            for (int j = 0; j < board[i].size(); j++) {
                if (this->board[i][j] < 0) return false;
            }
        }
        return true;
    }

    void computeInnerEdges() {
        for (const Cage& cage : cages) {
            for (const Position& p1 : cage.getPositions()) {
                for (const Position& p2 : cage.getPositions()) {
                    if (p1.sharesNorthEdge(p2)) innerEdges[p1].north = true;
                    if (p1.sharesSouthEdge(p2)) innerEdges[p1].south = true;
                    if (p1.sharesEastEdge(p2)) innerEdges[p1].east = true;
                    if (p1.sharesWestEdge(p2)) innerEdges[p1].west = true;
                }
            }
        }
    }

   public:
    Puzzle(std::vector<Cage> _cages, int _n)
        : cages(_cages), board(_n, std::vector<int>(_n, -1)) {
        // Validate puzzle definition
        for (int i = 0; i < cages.size(); ++i) {
            for (const Position& p : cages[i].getPositions()) {
                if (board[p.x][p.y] == 0) {
                    throw std::invalid_argument("Overlapping cages detected.");
                }
                board[p.x][p.y] = 0;
                cageMap[p] = &cages[i];
            }
        }
        if (!this->isInitialized()) {
            throw std::invalid_argument("Insufficient cage coverage detected.");
        }
        this->computeInnerEdges();
    }

    Board getBoard() const { return this->board; }

    void setBoard(Board b) { this->board = b; }

    bool columnHasValue(int x, int value) const {
        for (int i = 0; i < board.size(); ++i) {
            if (this->board[x][i] == value) {
                return true;
            }
        }
        return false;
    }

    bool rowHasValue(int y, int value) const {
        for (int i = 0; i < board.size(); ++i) {
            if (this->board[i][y] == value) {
                return true;
            }
        }
        return false;
    }

    const Cage& findCage(int x, int y) const {
        return *(this->cageMap.at(Position(x, y)));
    }

    std::vector<int> possibleValues(int x, int y) const {
        std::vector<int> results;
        if (this->board[x][y]) {
            // Early exit; this square is already filled
            return results;
        }
        for (int i = 1; i <= this->board.size(); ++i) {
            if (rowHasValue(y, i) || columnHasValue(x, i)) {
                continue;
            }
            const Cage& c = this->findCage(x, y);
            if (c.fits(this->board, i)) {
                results.push_back(i);
            }
        }
        return results;
    }

    bool isSolved() const {
        for (int i = 0; i < board.size(); ++i) {
            for (int j = 0; j < board[i].size(); j++) {
                if (this->board[i][j] <= 0) return false;
            }
        }
        return true;
    }

    void setValue(int x, int y, int value) { this->board[x][y] = value; }

    int getValue(int x, int y) const { return this->board[x][y]; }

    int size() const { return this->board.size(); }

    EdgeSet getEdgeSet(const Position& p) const {
        auto iEdge = this->innerEdges.find(p);
        if (iEdge != this->innerEdges.end()) {
            return iEdge->second;
        }
        EdgeSet defaultSet(true);
        if (p.x == this->size() && p.y >= 0) {
            defaultSet.west = false;
        }
        if (p.y == this->size()) {
            defaultSet.north = false;
            defaultSet.south = false;
        }
        if (p.y < 0 && p.x < this->size()) {
            defaultSet.south = false;
        }
        return defaultSet;
    }

    void printSep(int y, bool label) const {
        using namespace box;
        std::wcout << "  ";
        for (int x = 0; x <= this->size(); ++x) {
            bool printed = false;
            Position p(x, y);
            EdgeSet edges = getEdgeSet(p);
            wchar_t edge = edges.west ? N : V;
            if (label) {
                for (auto& cage : this->cages) {
                    if (p == cage.getRenderPosition()) {
                        std::wcout << edge << std::setw(5) << std::left << cage;
                        printed = true;
                    }
                }
            }
            if (!printed) std::wcout << edge << "     ";
        }
        std::wcout << std::endl;
    }

    void printDivide(int y, bool last) const {
        using namespace box;
        for (int x = 0; x < this->board.size(); ++x) {
            EdgeSet edges = getEdgeSet(Position(x, y));
            wchar_t edge = edges.north ? N : H;
            if (x == 0) {
                std::wcout << "  "
                           << (last ? UR
                                    : (y == 0 ? DR : (edges.north ? V : VR)));
            }
            std::wcout << edge << edge << edge << edge << edge;

            // Get neighboring edge set for corner computation
            EdgeSet neighbor = getEdgeSet(Position(x + 1, y - 1));
            EdgeSet corner(neighbor.west, edges.north, edges.east,
                           neighbor.south);
            std::wcout << corner.toCorner();
        }
        std::wcout << std::endl;
    }

    void prettyPrint() const {
        using namespace box;
        for (int y = 0; y < this->size(); ++y) {
            printDivide(y, false);
            printSep(y, true);
            std::wcout << "  " << V << "  ";
            for (int x = 0; x < this->size(); x++) {
                wchar_t edge = getEdgeSet(Position(x, y)).east ? N : V;
                std::wcout << this->board[x][y] << "  " << edge << "  ";
            }
            std::wcout << std::endl;
            printSep(y, false);
        }
        printDivide(this->size(), true);
    }
};

void solve(Puzzle& p) {
    Board prev = p.getBoard();
    for (int i = 0; i < p.size(); ++i) {
        for (int j = 0; j < p.size(); ++j) {
            if (p.getValue(i, j)) {
                continue;
            }
            auto possible = p.possibleValues(i, j);
            if (!possible.size()) {
                p.setBoard(prev);
                return;
            }
            for (auto v : possible) {
                p.setValue(i, j, v);
                solve(p);
                if (p.isSolved()) {
                    return;
                }
            }
        }
    }
    p.setBoard(prev);
}

int main(int argc, char** argv) {
    // Fucking locale bullshit
    constexpr char locale_name[] = "en_US.utf8";
    setlocale(LC_ALL, locale_name);
    std::locale::global(std::locale(locale_name));
    std::wcout.imbue(std::locale());

    // Define our puzzle, based on this image:
    // https://upload.wikimedia.org/wikipedia/commons/thumb/f/fd/KenKenProblem.svg/1024px-KenKenProblem.svg.png
    Puzzle p(
        {
            Cage({{0, 0}, {0, 1}}, 11, '+'),
            Cage({{1, 0}, {2, 0}}, 2, '/'),
            Cage({{3, 0}, {3, 1}}, 20, '*'),
            Cage({{4, 0}, {5, 0}, {5, 1}, {5, 2}}, 6, '*'),
            Cage({{1, 1}, {2, 1}}, 3, '-'),
            Cage({{4, 1}, {4, 2}}, 3, '/'),
            Cage({{0, 2}, {1, 2}, {0, 3}, {1, 3}}, 240, '*'),
            Cage({{2, 2}, {3, 2}}, 6, '*'),
            Cage({{0, 4}, {1, 4}}, 6, '*'),
            Cage({{2, 3}, {2, 4}}, 6, '*'),
            Cage({{3, 3}, {3, 4}, {4, 4}}, 7, '+'),
            Cage({{4, 3}, {5, 3}}, 30, '*'),
            Cage({{0, 5}, {1, 5}, {2, 5}}, 8, '+'),
            Cage({{3, 5}, {4, 5}}, 2, '/'),
            Cage({{5, 4}, {5, 5}}, 9, '+'),
        },
        6);

    // Solve the puzzle
    solve(p);
    if (!p.isSolved()) {
        std::wcout << "The puzzle could not be solved!" << std::endl;
        return 1;
    }
    p.prettyPrint();

    return 0;
}

Example output: