【前缀和 && 子矩阵的中位数】Gym - 101778F Median and Queries

Step1 Problem:

有一个 n * m 的矩阵，每个位置都有一个值 a[i][j]. 有 q 个询问。
对于每个询问给你一个子矩阵：
子矩阵的 左上角(x1, y1)，右下角(x2, y2)，求该子矩阵的中位数。
中位数：所有数排序后，中间的那一个。如果有两个中位数，选靠左边的那一个。
数据范围：
1<=n, m<=100, 1<=q<=1e5, 1<=a[i][j]<=500.

Step2 Ideas:

sum[i][j][k]：子矩阵 { (1, 1), (i, j) }，所有元素值小于等于 k 的个数。
我们可以利用前缀和求出 sum[i][j][k].
容斥求出子矩阵 { (x1, y1), (x2, y2) } 所有元素值小于等于 k 的个数：sum[x2][y2][k] - sum[x2][y1-1][k] - sum[x1-1][y2][k] + sum[x1-1][y1-1][k];

对于每个询问，我们知道子矩阵排序后第几个数是中位数。
我们可以二分答案 ans。

Step3 Code:

#include<bits/stdc++.h>
using namespace std;
int sum[105][105][505];
int a[105][105];
void init(int n, int m)
{
    for(int i = 1; i <= n; i++)
    {
        for(int j = 1; j <= m; j++)
        {
            for(int k = 1; k <= 500; k++)
                sum[i][j][k] = 0;
        }
    }
}
int main()
{
    int T, n, m, q;
    scanf("%d", &T);
    while(T--)
    {
        scanf("%d %d %d", &n, &m, &q);
        init(n, m);
        for(int i = 1; i <= n; i++)
        {
            for(int j = 1; j <= m; j++)
            {
                scanf("%d", &a[i][j]);
                sum[i][j][a[i][j]]++;
            }
        }
        for(int i = 1; i <= n; i++)
        {
            for(int j = 1; j <= m; j++)
            {
                for(int k = 1; k <= 500; k++)
                    sum[i][j][k] += sum[i][j-1][k];
            }
        }
        for(int i = 1; i <= n; i++)
        {
            for(int j = 1; j <= m; j++)
            {
                for(int k = 1; k <= 500; k++)
                    sum[i][j][k] += sum[i-1][j][k];
            }
        }
        for(int i = 1; i <= n; i++)
        {
            for(int j = 1; j <= m; j++)
            {
                for(int k = 1; k <= 500; k++)
                    sum[i][j][k] += sum[i][j][k-1];
            }
        }
        int x1, y1, x2, y2;
        while(q--)
        {
            scanf("%d %d %d %d", &x1, &y1, &x2, &y2);
            int total = (y2-y1+1)*(x2-x1+1);//总的个数
            int ave = (total+1)/2;//中位数是第几个
            int l = 1, r = 500, ans;
            while(l <= r) {
                int mid = (l+r)/2;
                int num = sum[x2][y2][mid] - sum[x2][y1-1][mid] - sum[x1-1][y2][mid] + sum[x1-1][y1-1][mid];
                if(num >= ave) {//不断逼近正确答案
                    ans = mid;
                    r = mid-1;
                }
                else l = mid+1;
            }
            printf("%d\n", ans);
        }
    }
    return 0;
}