LeetCode/problems/max-sum-of-rectangle-no-larger-than-k.cpp

#include <set>
#include <vector>

class Solution {
    void row_partial_sums(const std::vector<std::vector<int>>& matrix, std::size_t y, std::vector<int>& partial) const
    {
        for (std::size_t x = 0; x < partial.size(); x++) {
            partial[x] += matrix[y][x];
        }
    }

    int find(std::vector<int>& partial_sums, int k) const
    {
        int max_sum = INT_MIN;

        std::set<int> prefixes;
        prefixes.insert(0);

        int running_sum = 0;
        for (auto partial : partial_sums) {
            running_sum += partial;

            auto found = prefixes.lower_bound(running_sum - k);
            if (found != prefixes.end()) {
                max_sum = std::max(max_sum, running_sum - *found);
            }

            prefixes.insert(running_sum);
        }

        return max_sum;
    }

public:
    int maxSumSubmatrix(const std::vector<std::vector<int>>& matrix, int k) const
    {
        std::size_t rows = matrix.size();
        std::size_t cols = matrix.front().size();

        int max_sum = INT_MIN;

        for (std::size_t y_min = 0; y_min < rows; y_min++) {
            std::vector<int> partial_sums(cols, 0);

            for (std::size_t y_max = y_min; y_max < rows; y_max++) {
                row_partial_sums(matrix, y_max, partial_sums);
                int maximum = find(partial_sums, k);
                max_sum = std::max(max_sum, maximum);
            }
        }

        return max_sum;
    }
};