AtCoder Beginner Contest 320 E

https://atcoder.jp/contests/abc320/tasks/abc320_e

そうめんが流れてくるのと人が戻ってくるのをイベントとして、PriorityQueueに突っ込みます。ただ、そうめんが流れてくるイベントと人が戻ってくるイベントを別構造体とすると、なかなか大変です。結局、processという共通トレイトのメソッドの返る値で違いを吸収しています。

// Somen Nagashi
#![allow(non_snake_case)]

use std::collections::{BinaryHeap, 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()
}


//////////////////// Event ////////////////////

trait Event {
    fn time(&self) -> (u32, u32);
    fn process(&self) -> (usize, Option<&Somen>);
}


//////////////////// Somen ////////////////////

struct Somen {
    T: u32,
    W: u64,
    S: u32
}

impl Event for Somen {
    fn time(&self) -> (u32, u32) {
        (self.T, 1)
    }
    
    fn process(&self) -> (usize, Option<&Somen>) {
        (0, Some(&self))
    }
}

impl Somen {
    fn new(T: u32, W: u64, S: u32) -> Somen {
        Somen { T, W, S }
    }
    
    fn read() -> Somen {
        let v = read_vec();
        Somen::new(v[0], v[1] as u64, v[2])
    }
}


//////////////////// ComeBack ////////////////////

struct ComeBack {
    ID: usize,  // person's ID
    T: u32
}

impl Event for ComeBack {
    fn time(&self) -> (u32, u32) {
        (self.T, 0)
    }
    
    fn process(&self) -> (usize, Option<&Somen>) {
        (self.ID, None)
    }
}


//////////////////// EventWrapper ////////////////////

use std::cmp::Ordering;

pub struct EventWrapper(Box<dyn Event>);

impl PartialEq for EventWrapper {
    fn eq(&self, other: &Self) -> bool {
        self.0.time() == other.0.time()
    }
}

impl Eq for EventWrapper {}

impl PartialOrd for EventWrapper {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for EventWrapper {
    fn cmp(&self, other: &Self) -> Ordering {
        other.0.time().cmp(&self.0.time())
    }
}


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

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

fn F(N: usize, somens: Vec<Somen>) {
    let mut weights: Vec<u64> = vec![0; N];
    let mut line: BTreeSet<usize> = (0..N).collect();
    let mut events = BinaryHeap::new();
    for somen in somens.into_iter() {
        events.push(EventWrapper(Box::new(somen)))
    }
    
    while let Some(event_wrapper) = events.pop() {
        let event = event_wrapper.0;
        let (ID_CB, op) = event.process();
        match op {
            Some(somen) => {    // そうめんが流れてくる
                match line.pop_first() {
                    Some(ID) => {
                        weights[ID] += somen.W;
                        let c = ComeBack { ID, T: event.time().0 + somen.S };
                        events.push(EventWrapper(Box::new(c)))
                    },
                    None        => ()
                }
            },
            None => {           // 列に帰ってくる
                line.insert(ID_CB);
            }
        }
    }
    
    for w in weights.into_iter() {
        println!("{}", w)
    }
}

fn main() {
    let (N, somens) = read_input();
    F(N, somens)
}