LeetCode/cpp/matrix-diagonal-sum.cpp

#include <cassert>
#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;
        }

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

  public:
    int diagonalSum(const std::vector<std::vector<int>> &mat) {
        int sum = 0;

        indices down{0, 0}, up{0, static_cast<int>(mat.size()) - 1};
        indices d_down{1, 1}, d_up{1, -1};
        for (std::size_t i = 0; i < mat.size();
             ++i, down += d_down, up += d_up) {
            sum += down[mat];
            if (down != up) {
                sum += up[mat];
            }
        }

        return sum;
    }
};

int main() {
    Solution s;

    assert((s.diagonalSum({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}) == 25));
    assert(
        (s.diagonalSum(
             {{1, 1, 1, 1}, {1, 1, 1, 1}, {1, 1, 1, 1}, {1, 1, 1, 1}}) == 8));
    assert((s.diagonalSum({{5}}) == 5));
    return 0;
}