矩形分区选Top-K最大区域代码优化：解决超时问题

问题背景 #

我正在解决一个问题：将尺寸为w x h的矩形王国通过垂直和水平线条划分为多个区域，任务是选出k个面积最大的区域，并计算其中最大与最小面积的差值。

任务详情 #

2024年，T王国的奥列格国王迎来60岁生日，他关注遗产分配问题。当前T王国是w × h的矩形，被平行于边的线条划分为区域，垂直切割点坐标为x1, x2...xn，水平切割点坐标为y1, y2...ym。
例如，当w = 10, h = 8, x = [1, 5, 7], y = [4, 6]时，王国划分如截图所示。奥列格有k个儿子，他要给每个儿子分配一个区域，剩余区域由精英管理。他希望给儿子们分配k个面积最大的区域，长子得最大的，幼子得这k个里最小的，需要计算这两个区域的面积。

当前问题 #

代码通过了100个测试用例中的97个，但在第98个测试用例出现**运行时间超限（2秒）**的错误。

我的疑问 #

如何优化该算法以减少运行时间？能否避免双重循环或更高效地找出Top-K面积的区域？

当前代码 #

#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;
}

输入输出示例 #

优化建议 #

1. 移除不必要的PMR内存池 #

当前代码使用std::pmr内存池，实际高频操作中反而会增加内存管理开销。直接使用默认的std::priority_queue（基于std::vector）即可，减少复杂度。

2. 缩小visited数组的空间 #

原visited数组大小为(n+1)x(m+1)，但经过nth_element和resize后，widths和heights的最大长度都是k，因此visited只需k x k的大小，大幅减少内存占用和初始化时间。

3. 替换visited的实现为哈希集合 #

二维数组的访问可能存在缓存不友好问题，可将(i,j)编码为64位整数（例如static_cast<uint64_t>(i) << 32 | j），用std::unordered_set<uint64_t>记录已访问的坐标，提升访问速度。

4. 去掉冗余的排序步骤 #

std::nth_element已经保证前k个元素是最大的k个，后续的sort完全多余，可直接删除，节省O(k log k)的时间。修改后的宽度高度处理：

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);
}

5. 用结构体替代tuple提升访问效率 #

std::tuple的元素访问比结构体成员慢，可定义简单结构体存储数据：

struct AreaNode {
    long long area;
    int i;
    int j;
    AreaNode(long long a, int x, int y) : area(a), i(x), j(y) {}
    bool operator<(const AreaNode& other) const {
        return area < other.area;
    }
};

之后使用std::priority_queue<AreaNode>，减少模板开销，提升访问速度。

6. 确保切割点有序 #

如果输入的切割点坐标无序，需先对x、y数组排序，否则计算的宽度高度会出错，间接影响性能：

std::sort(x.begin() + 1, x.end() - 1);
std::sort(y.begin() + 1, y.end() - 1);

内容的提问来源于stack exchange，提问作者Rocowoco