diff --git a/problems/matrix-diagonal-sum.cpp b/problems/matrix-diagonal-sum.cpp new file mode 100644 index 0000000..5433d8c --- /dev/null +++ b/problems/matrix-diagonal-sum.cpp @@ -0,0 +1,46 @@ +#include +#include + +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>& mat) const { + return mat[y][x]; + } + }; + + public: + int diagonalSum(const std::vector>& mat) { + int sum = 0; + + indices down{0, 0}, up{0, static_cast(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; +}