LeetCode/problems/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;
}