Module d

Expand description

§ABC 413 D - Make Geometric Sequence

refs: https://atcoder.jp/contests/abc413/tasks/abc413_d

等比数列の比を求めようとするのは浮動小数点の誤差問題があるので, 等比になっているかどうかの判定は乗算にて行う. 一部例外を除き, 絶対値をキーとしたソートを実施した後に, 各要素間の比が同一になっているかを検証すれば良い.

このとき, 比が 1 の場合は, 全ての要素が同一である必要がある. また, 比が -1 の場合は, 負数と正数の数の差が 1 以下である必要がある.

要素数が 2 以下の場合は, 比の検証を行う必要はなく, 常に等比数列となる.

use std::collections::HashSet;

use proconio::input;

fn main() {
    input! { t: usize }

    for _ in 0..t {
        solve();
    }
}

fn solve() {
    input! {
        n: usize,
        mut a: [isize; n],
    }

    if n < 3 {
        println!("Yes");
        return;
    }

    if is_geometric_1(&a) || is_geometric_neg_1(&a) {
        println!("Yes");
        return;
    }

    a.sort_by_key(|a| a.abs());

    for i in 2..n {
        if a[i - 1] * a[i - 1] != a[i] * a[i - 2] {
            println!("No");
            return;
        }
    }

    println!("Yes");
}

fn is_geometric_1(a: &[isize]) -> bool {
    HashSet::<isize>::from_iter(a.iter().cloned()).len() == 1
}

fn is_geometric_neg_1(a: &[isize]) -> bool {
    let f = a[0];
    let fc = a.iter().filter(|&&x| x == f).count();
    let fnc = a.iter().filter(|&&x| x == -f).count();

    fc + fnc == a.len() && fc.abs_diff(fnc) <= 1
}