AtCoder Beginner Contest 439 D

https://atcoder.jp/contests/abc439/tasks/abc439_d

1個7に当たる値を決めれば5と3も決まるので、結局同じ値があったときが問題になります。jを固定して、jが最大と最小に分けて、最大ならそうなるギリギリのiとkを求めれば、iまでの個数とkまでの個数を掛け合わせたのが個数になります。

// Kadomatsu Subsequence
#![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()
}


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

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

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

fn collect(A: &Vec<i32>) -> HashMap<i32, Vec<usize>> {
    let mut m: HashMap<i32, Vec<usize>> = HashMap::new();
    for (i, &a) in A.iter().enumerate() {
        let e = m.entry(a).or_insert(vec![]);
        (*e).push(i)
    }
    m
}

fn F_each(u: &Vec<usize>, v: &Vec<usize>, w: &Vec<usize>) -> usize {
    let (L1, L2, L3) = (u.len(), v.len(), w.len());
    if L1 == 0 || L2 == 0 || L3 == 0 {
        return 0
    }
    
    let mut counter: usize = 0;
    // jがmax
    let mut p1: usize = 0;
    let mut p3: usize = 0;
    for p2 in 0..L2 {
        // p1とp3がギリギリjがmaxになるところを探す
        while p1 < L1 && u[p1] <= v[p2] {
            p1 += 1
        }
        while p3 < L3 && w[p3] <= v[p2] {
            p3 += 1
        }
        if p1 == 0 || p3 == 0 {
            continue
        }
        p1 -= 1;
        p3 -= 1;
        counter += (p1+1) * (p3+1)
    }
    
    // jがmin
    (p1, p3) = (0, 0);
    for p2 in 0..L2 {
        while p1 < L1 && u[p1] <= v[p2] {
            p1 += 1
        }
        while p3 < L3 && w[p3] <= v[p2] {
            p3 += 1
        }
        if p1 == L1 || p3 == L3 {
            break
        }
        counter += (L1-p1) * (L3-p3)
    }
    
    counter
}

fn F(A: Vec<i32>) -> usize {
    let map = collect(&A);
    let mut counter: usize = 0;
    let set: HashSet<i32> = A.into_iter().collect();
    for a in set {
        if a % 7 != 0 {
            continue
        }
        let b = a / 7 * 5;
        let c = a / 7 * 3;
        if let Some(u) = map.get(&a) {
            if let Some(v) = map.get(&b) {
                if let Some(w) = map.get(&c) {
                    counter += F_each(u, v, w)
                }
            }
        }
    }
    counter
}

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