How to optimize a code for rectangle partitioning to avoid time limit?

Problem: I am solving a problem where a rectangular kingdom of size w x h is divided into counties by vertical and horizontal lines. The task is to select the k largest counties and compute the difference in area between the largest and the smallest among them.

Task: In the year 2024, King Oleg of the T-Kingdom is celebrating his sixtieth birthday. Consequently, he is interested in the matter of inheritance. Currently, the T-Kingdom is a rectangle of size w × h, divided into counties parallel to its sides. The vertical cuts occur at coordinates x1, x2. . . xn, and the horizontal cuts at y1, y2. . . ym.

For example, if w = 10, h = 8, x = [1, 5, 7], and y = [4, 6], the T-Kingdom would look like the provided screenshot. Oleg has k sons, and he wishes to bequeath one county to each of them, while the remaining areas will be managed by the elite. Naturally, as any father would, Oleg wants to give his sons the k largest counties by area, with the eldest son receiving the largest county and the youngest son receiving the smallest county among the k. Now, Oleg is interested in finding the difference in area between the county given to the youngest son and the county given to the eldest son. Help Oleg calculate the area of the county given to the youngest son and the area of the county given to the eldest son.

Error now: The code passes the first 97 tests out of 100, but on the 98th test I get an error that the maximum operating time has been exceeded (2 seconds).

Question: How can this algorithm be optimized to reduce the running time? Perhaps there are ways to avoid the double loop or find the top-K counties by area more efficiently?

Current code (Edited):

#include <iostream>
#include <vector>
#include <queue>
#include <tuple>
#include <algorithm>
#include <climits>
#include <memory_resource>

struct compare {
    bool operator()(const std::tuple<long long, int, int>& a, const std::tuple<long long, int, int>& b) {
        return std::get<0>(a) < std::get<0>(b);
    }
};

int main() {
    long long w, h;
    std::cin >> w >> h;
    int n, m, k;
    std::cin >> n >> m >> k;

    std::vector<long long> x(n + 2), y(m + 2);
    x[0] = 0; x[n + 1] = w;
    y[0] = 0; y[m + 1] = h;

    for (int i = 1; i <= n; i++) std::cin >> x[i];
    for (int i = 1; i <= m; i++) std::cin >> y[i];

    std::vector<long long> widths(n + 1), heights(m + 1);
    for (int i = 1; i <= n + 1; i++) widths[i - 1] = x[i] - x[i - 1];
    for (int i = 1; i <= m + 1; i++) heights[i - 1] = y[i] - y[i - 1];

    if (widths.size() > k) {
        std::nth_element(widths.begin(), widths.begin() + k - 1, widths.end(), std::greater{});
        widths.resize(k);
    }
    if (heights.size() > k) {
        std::nth_element(heights.begin(), heights.begin() + k - 1, heights.end(), std::greater{});
        heights.resize(k);
    }

    std::sort(widths.begin(), widths.end(), std::greater{});
    std::sort(heights.begin(), heights.end(), std::greater{});

    std::pmr::unsynchronized_pool_resource pool;
    using queue_container_t = std::pmr::vector<std::tuple<long long, int, int>>;
    std::priority_queue<std::tuple<long long, int, int>, queue_container_t, compare> maxHeap{ std::pmr::polymorphic_allocator<queue_container_t>{&pool} };

    maxHeap.emplace(widths[0] * heights[0], 0, 0);

    std::vector<std::vector<bool>> visited(n + 1, std::vector<bool>(m + 1, false));
    visited[0][0] = true;

    long long smallest = LLONG_MAX, largest = LLONG_MIN;

    for (int count = 0; count < k; ++count) {
        auto [area, i, j] = maxHeap.top();
        maxHeap.pop();

        smallest = std::min(smallest, area);
        largest = std::max(largest, area);

        if (i + 1 < widths.size() && !visited[i + 1][j]) {
            maxHeap.emplace(widths[i + 1] * heights[j], i + 1, j);
            visited[i + 1][j] = true;
        }
        if (j + 1 < heights.size() && !visited[i][j + 1]) {
            maxHeap.emplace(widths[i] * heights[j + 1], i, j + 1);
            visited[i][j + 1] = true;
        }
    }

    std::cout << smallest << " " << largest << std::endl;
    return 0;
}

}

error kindom

I'm trying to calculate the area of all counties and then use a priority queue to find the k largest areas. (I was calculating all areas before). But still, the data seems to be very large in test 98. Any advice on how to optimize this?

Input	Output
10 8	6 16
3 2
6
1 5 7
4 6

Solution

you don't need to sort x and y, they are already sorted
you only need the Kth greatest heights and widths, you can use std::nth_element to do this operation, this will speed up the initial sort and reduce strain on memory at the next stage.

if (widths.size() > k)
{
    std::nth_element(widths.begin(), widths.begin() + k - 1, 
        widths.end(), std::greater{});
    widths.resize(k);
}

if (heights.size() > k)
{
    std::nth_element(heights.begin(), heights.begin() + k - 1, 
        heights.end(), std::greater{});
    heights.resize(k);
}

std::sort(widths.begin(), widths.end(), std::greater{});
std::sort(heights.begin(), heights.end(), std::greater{});

you actually need less than k elements for width and height, maybe you can use a better heuristic to reduce the size even more, but i can't come up with a better heuristic right now.

it might be worth looking into std::pmr::vector as the underlying container for the priority queue, with an unsynchronized_pool_resource to reduce the allocation and deallocation cost, as currently it is allocating and deallocating a lot of small vectors, which might be costly.

std::pmr::unsynchronized_pool_resource pool;
using queue_container_t = std::pmr::vector<std::tuple<long long, int, int>>;
std::priority_queue<std::tuple<long long, int, int>, queue_container_t, compare> maxHeap{ std::pmr::polymorphic_allocator<queue_container_t>{&pool} };

you can do this K skip, the basic idea is that we find the smallest rectangle that can be omitted from the search space whose elements count is less than K, once we remove it we only search the remaining strips which are way smaller.

if (k > 4)
{
    int current_x = 0;
    int current_y = 0;
    while (current_x + 2 < widths.size() && current_y + 2 < heights.size())
    {
        long long current_area_left = widths[current_x] * heights[current_y + 1];
        long long current_area_right = widths[current_x + 1] * heights[current_y];
        int new_current_y = current_y;
        int new_current_x = current_x;
        if (current_area_left > current_area_right)
        {
            new_current_y++;
        }
        else
        {
            new_current_x++;
        }
        int taken_k = (new_current_x + 1) * (new_current_y + 1);
        if (taken_k > k)
        {
            break;
        }
        current_x = new_current_x;
        current_y = new_current_y;
    }
    k -= (current_x + 1) * (current_y + 1);
    maxHeap.push(std::make_tuple(widths[current_x + 1] * heights[0], current_x + 1, 0));
    maxHeap.push(std::make_tuple(widths[0] * heights[current_y + 1], 0, current_y + 1));
}
else
{
    maxHeap.push(std::make_tuple(widths[0] * heights[0], 0, 0));
}


std::vector<std::vector<bool>> visited(n + 1, std::vector<bool>(m + 1, false));
visited[0][0] = true;

long long smallest = LLONG_MAX, largest = LLONG_MIN;
largest = widths[0] * heights[0];

note: this code replaces maxHeap.push(std::make_tuple(widths[0] * heights[0], 0, 0)); and you need to remove largest = std::max(largest, area); as it is pointless.