Module d 

ABC 342 D - Square Pair

refs: https://atcoder.jp/contests/abc342/tasks/abc342_d

平方数は素因数分解した際の因数の乗数がすべて偶数になる数. 乗数が奇数の因数だけが残るように事前処理しておくと, この処理後の値が同一 = 乗数が奇数となる因数が完全に合致する数なので, 乗算後の値が平方数になる.

処理後の数が同一になるものが N 個の場合, そこから 2 つ選択する組み合わせは $N \times (N - 2) / 2 $ で算出できる. 値が 0 である場合のみ, どの数と乗算しても平方数になるため例外処理が必要.

use std::collections::HashMap;

use proconio::input;

const MX: usize = 2 * 1e+5 as usize;

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

    let sq = square_num();

    for ai in a.iter_mut() {
        for &si in sq.iter() {
            if *ai % si == 0 {
                *ai /= si;
            }
        }
    }

    let mut map = HashMap::new();

    for ai in a {
        *map.entry(ai).or_insert(0usize) += 1;
    }

    let mut ans = 0;

    for (k, v) in map {
        if k == 0 {
            ans += v * (n - v);
        }
        ans += v * (v - 1) / 2;
    }

    println!("{ans}");
}

fn square_num() -> Vec<usize> {
    let mut result = vec![];
    let mut buf = 1usize;

    while result.last().is_none() || result.last() < Some(&MX) {
        result.push(buf.pow(2));
        buf += 1;
    }

    result.reverse();
    result
}