https://atcoder.jp/contests/abc440/tasks/abc440_e
多次元で途方もない量のデータを大きいほうから並べるという問題なので、PriorityQueueを使うのだろうとわかります。
Aを降順にソートして、Aに対応するクッキーの数を[K, 0, ...]とするとおいしさの和が最も大きくなります。これを初期状態として、ある状態からどの状態に移るかを決めれば終わりです。なるべく和が小さくならない方がいいので右隣に一つ移動すればよいです。例えば、[K-1, 1, 0, ...]なら[K-2, 2, 0, ...]と[K-1, 0, 1, ...]があります。
ただ、PriorityQueueを使うときにいつも問題になるのですが、辿っていくと同じ状態になることがあります。例えば、[K-2, 1, 1, 0, ...]は[K-2, 2, 0, 0, ...]と[K-1, 0, 1, 0, ...]から移動することができます。これを避けるように移動方法を決めなければなりません。こういうときはいつも逆向きに考えます。すなわち、[K-2, 1, 1, 0, ...]から逆向きの移動方法を一意に決めればよいです。こういうとき和の差が一番小さくなるように選ぶとよいのですが、今回の場合はそういうわけにはいきません。なので、決め打ちで一番右が移動するとします。そうすると、[K-2, 2, 0, 0, ...]から移るということになります。0でない一番左からかその隣からしか移動できないとすれば重複がなくなります。
// Cookies #![allow(non_snake_case)] use std::collections::BinaryHeap; //////////////////// 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() } //////////////////// Cookies //////////////////// #[derive(Debug, Eq, PartialEq)] struct Cookies { ns: Vec<i64>, s: i64 } impl Cookies { fn nexts(&self, A: &Vec<i64>) -> Vec<Cookies> { let mut cs: Vec<Cookies> = vec![]; let L = self.ns.len(); if A.len() > L { let mut ns: Vec<i64> = self.ns.clone(); ns[L-1] -= 1; ns.push(1); let s: i64 = self.s - A[L-1] + A[L]; cs.push(Cookies { ns, s }) } if L >= 2 && self.ns[L-2] > 0 { let mut ns: Vec<i64> = self.ns.clone(); ns[L-2] -= 1; ns[L-1] += 1; let s: i64 = self.s - A[L-2] + A[L-1]; cs.push(Cookies { ns, s }) } cs } fn create(K: i64, A: &Vec<i64>) -> Cookies { let ns: Vec<i64> = vec![K]; let s = A[0] * K; Cookies { ns, s } } } impl Ord for Cookies { fn cmp(&self, other: &Self) -> std::cmp::Ordering { self.s.cmp(&other.s) } } impl PartialOrd for Cookies { fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { Some(self.cmp(other)) } } //////////////////// process //////////////////// fn read_input() -> (i64, usize, Vec<i64>) { let v: Vec<usize> = read_vec(); let (K, X) = (v[1] as i64, v[2]); let A: Vec<i64> = read_vec(); (K, X, A) } fn F(K: i64, X: usize, mut A: Vec<i64>) { A.sort_by(|a, b| b.cmp(a)); let mut pq: BinaryHeap::<Cookies> = BinaryHeap::new(); let c0 = Cookies::create(K, &A); pq.push(c0); for _ in 0..X { let c = pq.pop().unwrap(); println!("{}", c.s); for c1 in c.nexts(&A) { pq.push(c1) } } } fn main() { let (K, X, A) = read_input(); F(K, X, A) }