leetcode/Problems/0498.cpp at master · DimitrijeDobrota/leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class Solution {
  public:
    bool valid(int i, int m, int j, int n) { return i >= 0 && i <= m && j >= 0 && j <= n; }

    void quick_adjust(int &i, int &j, bool &up) {
        if (up)
            i++;
        else
            j++;
        up = !up;
    }

    void move(int &i, int &j, bool &up) {
        if (up) {
            i--;
            j++;
        } else {
            i++;
            j--;
        }
    }

    vector<int> findDiagonalOrder(vector<vector<int>> &mat) {
        vector<int> res;
        bool up = true;
        int i = 0, j = 0;
        int m = mat.size() - 1, n = mat[0].size() - 1;

        while (true) {
            res.push_back(mat[i][j]);
            if (i == m && j == n) break;

            move(i, j, up);

            if (!valid(i, m, j, n)) {
                quick_adjust(i, j, up);
                while (!valid(i, m, j, n))
                    move(i, j, up);
            }
        }

        return res;
    }
};