G - Make Geometric Sequence 解説 by ngtkana


\(A\) を並び替えて等比数列 \(B\) が作れる条件を特定しましょう。

\(A\) の各要素の絶対値を取り、ソートしたものを \(B'\) と置きます。 まず、\(B'\) が等比数列でなければ「不可能」なので、これは等比数列だとしましょう。

このとき \(B\) の候補として、次の 4 つを試せば十分です(ただし \(\circ\) は各点積):

  1. \(B = B'\)
  2. \(B = -B'\)
  3. \(B = B' \circ (1, -1, 1, -1, \dots)\)
  4. \(B = -B' \circ (1, -1, 1, -1, \dots)\)

従って、実際にこれらを構築して、元の \(A\) の順列になっているかをチェックすればよいです。

コード

Rust 37 ms

use itertools::iproduct;
use itertools::Itertools;
use proconio::input;

fn main() {
    input! {
        q: usize,
    }
    for _ in 0..q {
        input! {
            n: usize,
            mut a: [i64; n],
        }
        a.sort_unstable();
        let mut b = a.iter().map(|&x| x.abs()).collect::<Vec<_>>();
        b.sort_unstable();
        let ans = b.iter().tuple_windows().all(|(x, y, z)| x * z == y * y)
            && iproduct!([1, -1], [1, -1]).any(|(s, t)| {
                let mut b = b
                    .iter()
                    .zip(std::iter::successors(Some(1), |&x| Some(t * x)))
                    .map(|(x, t)| s * t * x)
                    .collect::<Vec<_>>();
                b.sort_unstable();
                a == b
            });
        println!("{}", if ans { "Yes" } else { "No" });
    }
}

投稿日時:
最終更新: