G - Make Geometric Sequence 解説
by
ngtkana
\(A\) を並び替えて等比数列 \(B\) が作れる条件を特定しましょう。
\(A\) の各要素の絶対値を取り、ソートしたものを \(B'\) と置きます。 まず、\(B'\) が等比数列でなければ「不可能」なので、これは等比数列だとしましょう。
このとき \(B\) の候補として、次の 4 つを試せば十分です(ただし \(\circ\) は各点積):
- \(B = B'\)
- \(B = -B'\)
- \(B = B' \circ (1, -1, 1, -1, \dots)\)
- \(B = -B' \circ (1, -1, 1, -1, \dots)\)
従って、実際にこれらを構築して、元の \(A\) の順列になっているかをチェックすればよいです。
コード
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" });
}
}
投稿日時:
最終更新: