C++로 배우는 행렬 대각선 순회(Diagonal Traverse) 알고리즘

문제 소개

M × N 크기의 행렬이 주어졌을 때, 모든 요소를 대각선 순서로 탐색해야 하는 문제입니다. 이때 순회는 지그재그(zigzag) 패턴으로 진행되며, 대각선마다 이동 방향이 번갈아 바뀝니다. 예를 들어 다음과 같은 3 × 3 행렬이 있다고 가정해 보겠습니다.

이 행렬을 대각선 순서로 탐색한 결과는 [1, 2, 4, 7, 5, 3, 6, 8, 9]가 됩니다. 1에서 시작해 오른쪽 위 방향으로 이동했다가, 다음 대각선에서는 왼쪽 아래 방향으로 진행하는 식으로 방향이 교차하며 순회가 이루어집니다.

알고리즘 접근 방법

핵심 아이디어는 각 대각선에 속한 요소들을 임시 배열에 모은 뒤, 현재 진행 방향에 따라 필요하면 배열을 뒤집어 결과에 이어 붙이는 것입니다. 단계별 풀이 과정은 다음과 같습니다.

• 결과 배열 ret을 준비하고, 행 개수 n, 열 개수 m, 방향 플래그 down = false로 초기화합니다.
• 1단계 — 첫 번째 열에서 시작하는 대각선 처리: i를 0부터 n−1까지 반복하면서 시작점을 (x=i, y=0)으로 설정합니다. 임시 배열 temp를 만들고, x ≥ 0이고 y < m인 동안 matrix[x][y]를 temp에 추가한 뒤 x는 1 감소, y는 1 증가시켜 오른쪽 위 방향으로 이동합니다. 대각선 수집이 끝나면 down이 true일 경우 temp를 뒤집고, 모든 요소를 ret에 추가한 후 down 값을 반전시킵니다.
• 2단계 — 마지막 행에서 시작하는 대각선 처리: i를 1부터 m−1까지 반복하면서 시작점을 (x=n−1, y=i)로 설정합니다. 앞서와 동일한 방식으로 대각선 요소를 수집하고, 필요 시 뒤집어 ret에 추가한 뒤 down 값을 반전시킵니다.
• 모든 대각선을 처리한 후 ret을 반환합니다.

C++ 구현 코드

다음은 위 알고리즘을 C++로 구현한 전체 코드입니다.

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<int> v){
   cout << "[";
   for(int i = 0; i<v.size(); i++){
      cout << v[i] << ", ";
   }
   cout << "]"<<endl;
}
class Solution {
   public:
   vector<int> findDiagonalOrder(vector<vector<int>>& matrix) {
      vector <int> ret;
      int row = 0;
      int col = 0;
      int n = matrix.size();
      int m = n? matrix[0].size() : 0;
      bool down = false;
      for(int i = 0; i < n; i++){
         int x = i;
         int y = 0;
         vector <int> temp;
         while(x >= 0 && y < m){
            temp.push_back(matrix[x][y]);
            x--;
            y++;
         }
         if(down) reverse(temp.begin(), temp.end());
         for(int i = 0; i < temp.size(); i++)ret.push_back(temp[i]);
         down = !down;
      }
      for(int i = 1; i < m; i++){
         int x = n - 1;
         int y = i;
         vector <int> temp;
         while(x >= 0 && y < m){
            temp.push_back(matrix[x][y]);
            x--;
            y++;
         }
         if(down) reverse(temp.begin(), temp.end());
         for(int i = 0; i < temp.size(); i++)ret.push_back(temp[i]);
         down = !down;
      }
      return ret;
   }
};
main(){
   vector<vector<int>> v = {{1,2,3},{4,5,6},{7,8,9}};
   Solution ob;
   print_vector(ob.findDiagonalOrder(v));
}

실행 결과 확인

입력

[[1,2,3],[4,5,6],[7,8,9]]

출력

[1, 2, 4, 7, 5, 3, 6, 8, 9]

복잡도 분석

시간 복잡도: O(n × m) — 행렬의 모든 요소를 정확히 한 번씩 방문하므로 입력 크기에 비례합니다.
공간 복잡도: O(n × m) — 최종 결과를 담을 배열이 필요하며, 각 대각선을 임시로 저장하는 배열도 추가로 사용됩니다.