Submission #1482985

---

Source Code Expand

fn get_line() -> String {
    let mut s = String::new();
    std::io::stdin().read_line(&mut s).ok();
    s.trim().to_string()
}

fn reads<T>() -> Vec<T>
    where
        T: std::str::FromStr,
        <T as std::str::FromStr>::Err: std::fmt::Debug {
    get_line().split(' ').map(|x| x.parse().unwrap()).collect()
}

use std::collections::BinaryHeap;
use std::cmp::Ordering;

#[derive(Clone, Copy, Eq, PartialEq)]
struct E {
    num: u64,
    idx: usize,
}

impl Ord for E {
    fn cmp(&self, other: &Self) -> Ordering {
        self.num.cmp(&other.num)
    }
}

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

fn main() {
    let kt = reads();
    let xs = reads();

    let mut q = BinaryHeap::new();
    for (i, x) in xs.into_iter().enumerate() {
        let e = E { num: x, idx: i};
        q.push(e);
    }

    let mut res = 0;
    let mut prev = kt[1];
    while let Some(e) = q.pop() {
        if e.idx == prev {
            if let Some(next) = q.pop() {
                q.push(e);
                prev = next.idx;

                if next.num > 1 {
                    q.push( E {num: next.num - 1, idx: next.idx});
                }
            }
            else {
                res += e.num as u64;
            }
        }
        else {
            prev = e.idx;

            if e.num > 1 {
                q.push(E { num: e.num - 1, idx: e.idx});
            }
        }
    }

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

Submission Info

Judge Result