AtCoder Beginner Contest 366 E

https://atcoder.jp/contests/abc366/tasks/abc366_e

これも累積和を使います。
x方向とy方向は分解できるので、それぞれで計算して、あとで統合します。


はxとなる座標の個数
とすると、
xでのx方向の距離の和は、
なら、

なら、

なら、

となります。
x方向とy方向のそれぞれの距離の和が求められれば、あとは尺取り法的にDを超えない距離の和を求めればよいです。

// Manhattan Multifocal Ellipse
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// process ////////////////////

fn read_input() -> (usize, Vec<(i64, i64)>) {
    let v: Vec<usize> = read_vec();
    let (N, D) = (v[0], v[1]);
    let points: Vec<(i64, i64)> = (0..N).map(|_| read_vec::<i64>()).
                                            map(|v| (v[0], v[1])).collect();
    (D, points)
}

// xもyも最小が1にする
fn convert_points(points: Vec<(i64, i64)>) -> Vec<(usize, usize)> {
    let min_x = points.iter().map(|&(x, _)| x).min().unwrap();
    let min_y = points.iter().map(|&(_, y)| y).min().unwrap();
    points.into_iter().map(|(x, y)| ((x - min_x + 1) as usize,
                                     (y - min_y + 1) as usize)).
                                                collect::<Vec<(usize, usize)>>()
}

fn accumulate(xs: Vec<usize>) -> (Vec<usize>, Vec<usize>) {
    let min_x = *xs.iter().next().unwrap();
    let max_x = *xs.iter().last().unwrap();
    let mut ns: Vec<usize> = vec![0; max_x - min_x + 1];
    for x in xs.into_iter() {
        ns[x-min_x] += 1
    }
    let mut a: Vec<usize> = vec![0; max_x - min_x + 2];
    let mut c: Vec<usize> = vec![0; max_x - min_x + 2];
    for (i, n) in ns.into_iter().enumerate() {
        a[i+1] += a[i] + n * (min_x + i);
        c[i+1] += c[i] + n
    }
    (a, c)
}

fn sum_dist(mut xs: Vec<usize>) -> (Vec<usize>, Vec<usize>) {
    xs.sort();
    let (a, c) = accumulate(xs);
    let L = a.len();
    let last_a = a[L-1];
    let last_c = c[L-1];
    let mut dists = vec![0; L];
    for i in 0..a.len() {
        dists[i] = i * 2 * c[i] + last_a - i * last_c - 2 * a[i];
    }
    
    // 個数が半分のところで分割して両方とも昇順になるようにする
    let mid = (0..L).filter(|&i| c[i] * 2 >= last_c).next().unwrap();
    (dists[..mid].into_iter().rev().map(|&x| x).collect::<Vec<usize>>(),
        dists[mid..].to_vec())
}

fn count_points(ds1: &Vec<usize>, ds2: &Vec<usize>,
                                N: usize, D: usize) -> usize {
    let dist = |l, ds: &Vec<usize>| {
        if l < ds.len() {
            ds[l]
        }
        else {
            ds[ds.len()-1] + (l - ds.len() + 1) * N
        }
    };
    
    let mut counter: usize = 0;
    if ds1[0] + ds2[0] > D {
        return 0
    }
    
    let mut l = 0;
    while dist(0, ds1) + dist(l, ds2) <= D {
        l += 1
    }
    l -= 1;
    counter += l + 1;
    for k in 1.. {
        if dist(k, ds1) + ds2[0] > D {
            break
        }
        else if l == 0 {
            counter += 1;
            continue
        }
        while dist(k, ds1) + dist(l, ds2) > D {
            l -= 1
        }
        counter += l + 1
    }
    counter
}

fn F(D: usize, points_: Vec<(i64, i64)>) -> usize {
    let N = points_.len();
    let points = convert_points(points_);
    let xs: Vec<usize> = points.iter().map(|&(x, _)| x).collect();
    let ys: Vec<usize> = points.iter().map(|&(_, y)| y).collect();
    let (dist_xs1, dist_xs2) = sum_dist(xs);
    let (dist_ys1, dist_ys2) = sum_dist(ys);
    count_points(&dist_xs1, &dist_ys1, N, D) +
    count_points(&dist_xs1, &dist_ys2, N, D) +
    count_points(&dist_xs2, &dist_ys1, N, D) +
    count_points(&dist_xs2, &dist_ys2, N, D)
}

fn main() {
    let (D, points) = read_input();
    println!("{}", F(D, points))
}