5. Sort Values by Bouncing Diagonal Sums

1. Description

You are given an n × n integer matrix.
For each cell in the leftmost column, follow a diagonal path that initially moves up-right. When the path reaches the top row, it bounces and continues down-right until reaching the rightmost column.
Define the weight of a leftmost-column cell as the sum of all values along its path.
Return the values in the leftmost column sorted by their weights in ascending order. If two values have the same weight, sort them by value in ascending order.

2. Solution

n = matrix.size()
Time complexity: O(n$^2$)
Space complexity: O(n)

 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
vector<int> solution(const vector<vector<int>>& matrix) {
    const int n = matrix.size();
    vector<pair<int, int>> value_sums;

    for (int i = 0; i < n; ++i) {
        int r = i, c = 0;
        int sum = 0;

        while (r > 0) {
            sum += matrix[r--][c++];
        }

        while (c < n) {
            sum += matrix[r++][c++];
        }

        value_sums.emplace_back(matrix[i][0], sum);
    }

    sort(value_sums.begin(), value_sums.end(),
         [](const auto& a, const auto& b) {
             if (a.second != b.second) {
                 return a.second < b.second;
             }
             return a.first < b.first;
         });

    vector<int> result;
    result.reserve(n);
    for (const auto& [value, sum] : value_sums) {
        result.push_back(value);
    }

    return result;
}
comments powered by Disqus