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: