D - Pond Editorial by lucifer1004


If the median of a section is \(M\), then there should be at least \(\left\lceil\dfrac{K^2}{2}\right\rceil\) numbers in this section that satisfy \(A_{ij}\le M\).

Based on this observation, we can binary search for the answer. Given the guess of the median \(M\), we first turn the original matrix into a \(0\)-\(1\) matrix, in which \(0\) means \(A_{ij}>M\) and \(1\) means \(A_{ij}\le M\). Then we can find the sum of all \(K\times K\) sections in \(\mathcal{O}(N^2)\) time based on the precalculated \(2\)-D prefix sum. If the maximum among all section sums is smaller than \(\left\lceil\dfrac{K^2}{2}\right\rceil\), it means the guessed median is too large, and vice versa.

  • Time complexity is \(\mathcal{O}(N^2\log\operatorname{AMAX})\), in which \(\operatorname{AMAX}=10^9\).
  • Space complexity is \(\mathcal{O}(N^2)\).
use proconio::input;

fn main() {
    input! {
        n: usize,
        k: usize,
        a: [[usize; n]; n],
    }

    let pos = (k * k - 1) / 2 + 1;

    let mut lo: i32 = 0;
    let mut hi: i32 = 1_000_000_000;
    let mut v = vec![vec![0; n]; n];
    let mut s = vec![vec![0; n + 1]; n + 1];

    while lo <= hi {
        let mid = (lo + hi) >> 1;
        for i in 0..n {
            for j in 0..n {
                if a[i][j] <= mid as usize {
                    v[i][j] = 1;
                } else {
                    v[i][j] = 0;
                }
            }
        }
        for i in 1..=n {
            for j in 1..=n {
                s[i][j] = s[i - 1][j] + s[i][j - 1] - s[i - 1][j - 1] + v[i - 1][j - 1];
            }
        }

        if s[n][n] < pos {
            lo = mid + 1;
            continue;
        }

        let mut found = false;
        for i in k..=n {
            for j in k..=n {
                if s[i][j] - s[i - k][j] - s[i][j - k] + s[i - k][j - k] >= pos {
                    found = true;
                    break;
                }
            }
            if found {
                break;
            }
        }

        if found {
            hi = mid - 1;
        } else {
            lo = mid + 1;
        }
    }

    println!("{}", lo);
}

posted:
last update: