AtCoder Beginner Contest 331 E

https://atcoder.jp/contests/abc331/tasks/abc331_e

PriorityQueueを使って、合計が多い組み合わせから列挙していきます。

// Set Meal
#![allow(non_snake_case)]

use std::collections::{BinaryHeap, HashSet};


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


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

fn read_bad_combination() -> (usize, usize) {
    let v = read_vec::<usize>();
    let (c, d) = (v[0]-1, v[1]-1);
    (c, d)
}

fn read_input() -> (Vec<i64>, Vec<i64>, Vec<(usize, usize)>) {
    let v = read_vec();
    let L: usize = v[2];
    let a: Vec<i64> = read_vec();
    let b: Vec<i64> = read_vec();
    let bad_combs = (0..L).map(|_| read_bad_combination()).collect::<Vec<_>>();
    (a, b, bad_combs)
}

type State = (i64, usize, usize);

fn nexts(s: State, A: &Vec<(usize, i64)>, B: &Vec<(usize, i64)>) -> Vec<State> {
    let mut neighbors: Vec<State> = vec![];
    let (_, i, j) = s;
    if i == 0 && j+1 < B.len() {
        neighbors.push((A[i].1 + B[j+1].1, i, j+1))
    }
    if i+1 < A.len() {
        neighbors.push((A[i+1].1 + B[j].1, i+1, j))
    }
    neighbors
}

fn F(a: Vec<i64>, b: Vec<i64>, bad_combs: Vec<(usize, usize)>) -> i64 {
    let mut A = a.iter().enumerate().map(|(i, &e)| (i, e)).collect::<Vec<_>>();
    let mut B = b.iter().enumerate().map(|(j, &e)| (j, e)).collect::<Vec<_>>();
    A.sort_by_key(|(_, e)| -e);
    B.sort_by_key(|(_, e)| -e);
    let set_bad_combs: HashSet<(usize, usize)> =
                            bad_combs.into_iter().collect();
    
    let mut heap = BinaryHeap::new();
    heap.push((A[0].1 + B[0].1, 0, 0));
    while let Some(s) = heap.pop() {
        let (price, i, j) = s;
        let comb = (A[i].0, B[j].0);
        if !set_bad_combs.contains(&comb) {
            return price
        }
        let neighbors = nexts(s, &A, &B);
        for s1 in neighbors.into_iter() {
            heap.push(s1)
        }
    }
    0
}

fn main() {
    let (a, b, bad_combs) = read_input();
    println!("{}", F(a, b, bad_combs))
}