AtCoder Beginner Contest 385 D

https://atcoder.jp/contests/abc385/tasks/abc385_d

同じ場所を何度も通ってそこにたくさんの点があると時間がかかるので、重複を除いてあらかじめ通る場所を計算しておきます。そのために、各固定座標に対し範囲をBTreeSetにし、そこに範囲を追加して、範囲が結合できるなら結合するようにします。

// Santa Claus 2
#![allow(non_snake_case)]

use std::collections::{HashMap, BTreeSet};


//////////////////// 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()
}


//////////////////// Point ////////////////////

type Point = (i64, i64);

fn read_point() -> Point {
    let v: Vec<i64> = read_vec();
    (v[0], v[1])
}


//////////////////// Range ////////////////////

type Range = (i64, i64);

fn add_range(rng: Range, ranges: &mut BTreeSet<Range>) {
    let mut to_remove = Vec::new();
    let mut start = rng.0;
    let mut end = rng.1;
    
    for &range in ranges.range((rng.0, i64::MIN)..(rng.1, i64::MAX)) {
        if range.1 < rng.0 || range.0 > rng.1 {
            continue
        }
        to_remove.push(range);
        start = start.min(range.0);
        end = end.max(range.1)
    }
    
    // 影響を受ける範囲を削除
    for range in to_remove {
        ranges.remove(&range);
    }
    
    // 新しいマージされた範囲を追加
    ranges.insert((start, end));

}


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

fn read_input() -> (usize, Point, Vec<Point>) {
    let v: Vec<i64> = read_vec();
    let (N, M) = (v[0] as usize, v[1] as usize);
    let S = (v[2], v[3]);
    let houses: Vec<Point> = (0..N).map(|_| read_point()).collect();
    (M, S, houses)
}

fn classify(keys: &Vec<i64>, values: &Vec<i64>) -> HashMap<i64, BTreeSet<i64>> {
    let mut c: HashMap<i64, BTreeSet<i64>> = HashMap::new();
    for (&k, &v) in keys.iter().zip(values.iter()) {
        let e = c.entry(k).or_insert(BTreeSet::new());
        (*e).insert(v);
    }
    c
}

fn read_move() -> (char, i64) {
    let v: Vec<String> = read_vec();
    let D: char = v[0].chars().next().unwrap();
    let C: i64 = v[1].parse().unwrap();
    (D, C)
}

fn collect_moves(moves: Vec<(char, i64)>, start: Point) ->
                                    (HashMap<i64, BTreeSet<Range>>,
                                     HashMap<i64, BTreeSet<Range>>, Point) {
    let mut range_x_dir: HashMap<i64, BTreeSet<Range>> = HashMap::new();
    let mut range_y_dir: HashMap<i64, BTreeSet<Range>> = HashMap::new();
    let mut x = start.0;
    let mut y = start.1;
    for (D, C) in moves.into_iter() {
        match D {
            'U' => {
                let e = range_y_dir.entry(x).or_insert(BTreeSet::new());
                add_range((y, y+C+1), &mut *e);
                y += C
            }
            'D' => {
                let e = range_y_dir.entry(x).or_insert(BTreeSet::new());
                add_range((y-C, y+1), &mut *e);
                y -= C
            }
            'L' => {
                let e = range_x_dir.entry(y).or_insert(BTreeSet::new());
                add_range((x-C, x+1), &mut *e);
                x -= C
            }
            _   => {
                let e = range_x_dir.entry(y).or_insert(BTreeSet::new());
                add_range((x, x+C+1), &mut *e);
                x += C
            }
        }
    }
    (range_x_dir, range_y_dir, (x, y))
}

fn count(set_xs: &BTreeSet<i64>, x_ranges: &BTreeSet<Range>) -> Vec<i64> {
    let xs: Vec<i64> = set_xs.iter().map(|&x| x).collect();
    let ranges: Vec<Range> = x_ranges.iter().map(|&r| r).collect();
    let mut v: Vec<i64> = vec![];
    let mut k: usize = 0;
    let mut l: usize = 0;
    while k < xs.len() && l < ranges.len() {
        let x = xs[k];
        let (first, last) = ranges[l];
        if first <= x && x < last {
            v.push(x);
            k += 1
        }
        else if x < first {
            k += 1
        }
        else {
            l += 1
        }
    }
    v
}

fn F(M: usize, S: Point, houses: Vec<Point>) {
    let xs: Vec<i64> = houses.iter().map(|&(x, _)| x).collect();
    let ys: Vec<i64> = houses.iter().map(|&(_, y)| y).collect();
    let collection_by_xs = classify(&xs, &ys);
    let collection_by_ys = classify(&ys, &xs);
    let moves: Vec<(char, i64)> = (0..M).map(|_| read_move()).collect();
    let (range_x_dir, range_y_dir, G) = collect_moves(moves, S);
    
    let mut set: BTreeSet<Point> = BTreeSet::new();
    for (y, x_ranges) in range_x_dir.into_iter() {
        match collection_by_ys.get(&y) {
            Some(set_xs) => {
                let xs1 = count(set_xs, &x_ranges);
                for x in xs1.into_iter() {
                    set.insert((x, y));
                }
            },
            None       => ()
        }
    }
    for (x, y_ranges) in range_y_dir.into_iter() {
        match collection_by_xs.get(&x) {
            Some(set_ys) => {
                let ys1 = count(set_ys, &y_ranges);
                for y in ys1.into_iter() {
                    set.insert((x, y));
                }
            },
            None       => ()
        }
    }
    println!("{} {} {}", G.0, G.1, set.len())
}

fn main() {
    let (M, S, houses) = read_input();
    F(M, S, houses)
}