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;
}