Given the shape of a ship and the location of one cell, bomb the remaining cells in minimum attempts

battleship.cpp

#include <iostream>
#include <map>
#include <random>
#include <tuple>
#include <vector>

std::random_device rd;
std::mt19937 gen(rd());

typedef std::vector<std::string> pattern;
typedef std::pair<int, int> coord;

class battleship {
  int height, width, size;
  int x0, y0, cnt_bombed;
  pattern ship, bombed;

 public:
  // Initialise battleship with size targets in a height x width area
  battleship(int height_, int width_, int size_)
      : height(height_), width(width_), size(size_) {
    ship = pattern(height, std::string(width, '.'));
    bombed = ship;
    cnt_bombed = 0;

    std::uniform_int_distribution<int> cell(0, height * width - 1);
    for (int i = 0; i < size; ++i) {
      int cur_cell = cell(gen);
      int x = cur_cell / width, y = cur_cell - x * width;
      if (ship[x][y] == '#') {
        --i;
        continue;
      }
      ship[x][y] = '#';
    }

    std::uniform_int_distribution<int> coord(-100, 100);
    x0 = coord(gen);
    y0 = coord(gen);
  }

  std::tuple<pattern, int, int> get_info() {
    int x, y;
    std::uniform_int_distribution<int> cell(0, height * width - 1);
    do {
      int cur_cell = cell(gen);
      x = cur_cell / width;
      y = cur_cell - x * width;
    } while (ship[x][y] == '.');
    return {ship, x + x0, y + y0};
  }

  int bomb(int x, int y) {
    int xx = x - x0, yy = y - y0;
    if (xx < 0 or yy < 0 or xx >= height or yy >= width) {
      std::cout << "Attack unsuccessful at (" << x << "," << y << ")"
                << std::endl;
      return 0;
    } else if (ship[xx][yy] == '.') {
      std::cout << "Attack unsuccessful at (" << x << "," << y << ")"
                << std::endl;
      return 0;
    } else if (bombed[xx][yy] == '#') {
      std::cout << "Attack unsuccessful on previously destroyed location at ("
                << x << "," << y << ")" << std::endl;
      return 0;
    } else {
      std::cout << "Attack successful at (" << x << "," << y << ")"
                << std::endl;
      bombed[xx][yy] = '#';
      if (++cnt_bombed == size) {
        std::cout << "Ship destroyed!" << std::endl;
      }
      return 1;
    }
  }
};

class player {
  pattern ship;
  int height, width, startx, starty;
  std::vector<coord> coord_vec;
  std::vector<std::vector<int>> poss;
  std::vector<unsigned long long> mask_vec;  // TODO: Replace with bitset

  // How many bad cells do you hit on average if you play optimally?
  std::vector<std::pair<int, double>> best_move;

 public:
  player(pattern ship_, int startx_, int starty_)
      : ship(ship_), startx(startx_), starty(starty_) {
    height = ship.size();
    width = ship[0].size();
  }

  void generate_masks() {
    coord_vec.clear();
    poss.clear();
    std::map<coord, int> coord_map;

    for (int i = 0; i < height * width; ++i) {
      int x = i / width, y = i - x * width;
      if (ship[x][y] == '.') continue;
      poss.push_back(std::vector<int>());

      for (int j = 0; j < height * width; ++j) {
        int xx = j / width, yy = j - xx * width;
        if (ship[xx][yy] == '.') continue;

        int dx = xx - x, dy = yy - y;
        auto mit = coord_map.find({dx, dy});
        if (mit == coord_map.end()) {
          coord_map[{dx, dy}] = coord_vec.size();
          poss.back().push_back(coord_vec.size());
          coord_vec.push_back({dx, dy});
        } else {
          poss.back().push_back(mit->second);
        }
      }
    }

    mask_vec.clear();
    for (const auto& vec : poss) {
      unsigned long long mask = 0;
      for (auto i : vec) {
        mask |= (1ull << i);
      }
      mask_vec.push_back(mask);
    }
  }

  std::pair<int, double> get_best_move(int subset_mask) {
    auto& ret = best_move[subset_mask];
    if (ret.first != -1) return ret;
    if ((subset_mask & (subset_mask - 1)) == 0) return ret = {-2, 0};

    ret = {-1, height * width};
    for (int i = 0; i < coord_vec.size(); ++i) {
      int subset_mask_0 = 0, subset_mask_1 = 0, cnt_mask_0 = 0, cnt_mask_1 = 0;
      for (int j = 0; j < poss.size(); ++j) {
        if ((subset_mask & (1 << j)) == 0) continue;
        if (mask_vec[j] & (1ull << i)) {
          subset_mask_1 |= (1 << j);
          cnt_mask_1++;
        } else {
          subset_mask_0 |= (1 << j);
          cnt_mask_0++;
        }
      }

      if (cnt_mask_0 == 0 or cnt_mask_1 == 0) continue;
      double cur_move =
          (cnt_mask_1 * get_best_move(subset_mask_1).second +
           cnt_mask_0 * (get_best_move(subset_mask_0).second + 1)) /
          (cnt_mask_0 + cnt_mask_1);
      if (cur_move < ret.second) ret = {i, cur_move};
    }
    return ret;
  }

  void initialise_DP() {
    best_move = std::vector<std::pair<int, double>>(1 << poss.size(), {-1, -1});
    get_best_move((1 << poss.size()) - 1);
  }

  int submask(int mask, int pos, unsigned long long val) {
    val <<= pos;
    for (int i = 0; i < poss.size(); ++i) {
      if ((mask & (1 << i)) == 0) continue;
      if ((mask_vec[i] & (1ull << pos)) != val) mask ^= (1 << i);
    }
    return mask;
  }

  void play(battleship& B) {
    unsigned long long tried = 0;
    int subset_mask = (1 << poss.size()) - 1;
    while (best_move[subset_mask].first != -2) {
      int move = best_move[subset_mask].first;
      tried |= (1ull << move);
      int ret = B.bomb(startx + coord_vec[move].first,
                       starty + coord_vec[move].second);
      subset_mask = submask(subset_mask, move, ret);
    }

    int id = -1;
    while (subset_mask) id++, subset_mask >>= 1;
    for (int i = 0; i < coord_vec.size(); ++i) {
      if ((tried & (1ull << i)) == 0 and (mask_vec[id] & (1ull << i))) {
        B.bomb(startx + coord_vec[i].first, starty + coord_vec[i].second);
      }
    }
  }
};

int main() {
  battleship B(2, 10, 10);
  pattern ship;
  int startx, starty;
  std::tie(ship, startx, starty) = B.get_info();
  std::cout << "Ship pattern:" << std::endl;
  for (const auto& row : ship) {
    std::cout << row << std::endl;
  }
  std::cout << "Start coordinates: (" << startx << "," << starty << ")"
            << std::endl;

  player P(ship, startx, starty);
  P.generate_masks();
  P.initialise_DP();
  P.play(B);
  return 0;
}