C++で行列を対角順に走査する方法（対角トラバーサル）

問題の概要

M×N の要素を持つ行列が与えられたとき、そのすべての要素を対角順（ジグザグ順）で取り出すことを考えます。例として、次のような 3×3 の行列を見てみましょう。

この場合、期待される出力は [1, 2, 4, 7, 5, 3, 6, 8, 9] です。要素は左上からスタートし、右上方向への斜め移動と左下方向への斜め移動を交互に繰り返しながら順番に辿られていきます。

アルゴリズムの手順

この問題は「右肩上がりの斜め列（反対角線）ごとに要素を集め、進む方向を交互に入れ替える」という方針で解くことができます。具体的な手順は以下の通りです。

• 結果格納用の配列 ret を用意し、n := 行数、m := 列数、down := false として初期化します。
• i を 0 ～ n-1 の範囲で繰り返します。
  • x := i、y := 0 とし、一時配列 temp を作成します。
  • x ≥ 0 かつ y ＜ m の間、matrix[x][y] を temp に追加し、x を 1 減らして y を 1 増やします。
  • down が true の場合は temp を反転します。
  • temp の全要素を ret に追加します。
  • down の真偽値を反転させます。
• 続けて i を 1 ～ m-1 の範囲で繰り返します。
  • x := n-1、y := i とし、一時配列 temp を作成します。
  • x ≥ 0 かつ y ＜ m の間、matrix[x][y] を temp に追加し、x を 1 減らして y を 1 増やします。
  • down が true の場合は temp を反転します。
  • temp の全要素を ret に追加します。
  • down の真偽値を反転させます。
• 最後に ret を返します。

最初のループでは左端の各セルから始まる斜め列を、2 番目のループでは最下行から始まる残りの斜め列を処理しています。これにより、行列全体を漏れなく一度ずつ走査できます。

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]