AtCoder Beginner Contest 359 F

https://atcoder.jp/contests/abc359/tasks/abc359_f

次数は1以上で全部で2N-2なので、まず全ての頂点に1を割り当てて、それからコストが小さいところから次数を1ずつ割り当てます。から次数を1つ割り当てると、コストが増すことになるので、それを考慮してPriorityQueueを使います。

// Tree Degree Optimization
#![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()
}


//////////////////// BinaryHeap ////////////////////

use std::collections::BinaryHeap;

#[derive(Debug, Eq, PartialEq)]
struct Item {
    value: i64,
    a: i64,
    d: i64
}

impl Ord for Item {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        other.value.cmp(&self.value)
    }
}

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


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

fn read_input() -> Vec<i64> {
    let _N: usize = read();
    let A: Vec<i64> = read_vec();
    A
}

fn F(A: Vec<i64>) -> i64 {
    let N = A.len();
    let mut heap = BinaryHeap::new();
    let mut s: i64 = 0;
    for a in A.into_iter() {
        s += a;
        heap.push(Item { value: a*3, a: a, d: 2 })
    }
    
    for _ in 0..N-2 {
        let Item { value, a, d } = heap.pop().unwrap();
        s += value;
        heap.push(Item { value: a*(d*2+1), a: a, d: d+1 })
    }
    s
}

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