LeetCode/cpp/spiral-matrix-ii.cpp

#include <cassert>
#include <iostream>
#include <vector>

class Solution {
  struct indices {
    int x;
    int y;

    bool operator==(const indices& other) const = default;
    indices& operator+=(const indices& other) {
      x += other.x;
      y += other.y;
      return *this;
    }

    friend indices operator+(indices left, const indices& right) {
      return left += right;
    }

    int& operator[](std::vector<std::vector<int>>& mat) const {
      return mat[y][x];
    }
  };

  struct spiral_indices {
    spiral_indices(const std::vector<std::vector<int>>& matrix)
        : x_bounds{-1, static_cast<int>(matrix.size() ? matrix[0].size() : 0)},
          y_bounds{-1, static_cast<int>(matrix.size())} {}

    bool done() const { return !in_bounds(idx); }

    spiral_indices& operator++() {
      // update bounds and change the direction if cannot move
      if (!in_bounds(idx + d)) {
        // change the direction
        d = {-d.y, d.x};

        // get the magnitude of the vector
        auto flat_d = d.x + d.y;

        // decide whether we're moving x or y bounds
        auto& bounds = d.x ? x_bounds : y_bounds;
        if (flat_d > 0) {
          bounds.x++;
        } else {
          bounds.y--;
        }
      }

      idx += d;

      return *this;
    }

    int& operator[](std::vector<std::vector<int>>& mat) const {
      return idx[mat];
    }

   private:
    indices idx{0, 0};
    indices d{1, 0};
    indices x_bounds, y_bounds;

    bool in_bounds(const indices& idx) const {
      return (x_bounds.x < idx.x && idx.x < x_bounds.y) &&
             (y_bounds.x < idx.y && idx.y < y_bounds.y);
    }
  };

 public:
  std::vector<std::vector<int>> generateMatrix(int n) {
    auto length = static_cast<std::size_t>(n);
    std::vector<std::vector<int>> spiral{length, std::vector<int>(length, 0)};

    int i = 1;
    for (spiral_indices idx{spiral}; !idx.done(); ++idx) {
      idx[spiral] = i++;
    }

    return spiral;
  }
};

int main() {
  Solution s;

  assert((s.generateMatrix(0) == std::vector<std::vector<int>>{}));
  assert((s.generateMatrix(1) == std::vector<std::vector<int>>{{1}}));
  assert(
      (s.generateMatrix(2) == std::vector<std::vector<int>>{{1, 2}, {4, 3}}));

  for (auto&& row : s.generateMatrix(3)) {
    for (auto x : row) std::cout << x << " ";
    std::cout << "\n";
  }

  assert((s.generateMatrix(3) ==
          std::vector<std::vector<int>>{{1, 2, 3}, {8, 9, 4}, {7, 6, 5}}));

  return 0;
}