問題

$n$ 個のバケツがあり、 $i$ 番目のものには1辺の長さが $1$ の正方形のブロックがある。

すべてのブロックを使用して、正方形を作成することができるか判定せよ。

入力

まず最初の1行目に、テストケースの個数を表す整数 $t$ が与えられる。

その後、 $t$ 個のテストケースのそれぞれに対して、以下のように入力が与えられる。

最初の1行目には、整数 $n$ が与えられる。
次の1行には、 $n$ 個の整数 ${a}_{1}, \cdots , {a}_{n}$ が順に与えられる。

条件

実行時間制限: 1s
メモリ制限: 256MB
$1 \leq t \leq {10}^{4}$
- すべてのテストケースにおける $n$ の値の合計は $2 \cdot {10}^{5}$ を超えない。
$1 \leq {a}_{i} \leq {10}^{9}$

出力

$t$ 行にわたり、各テストケースについて、正方形を作成できるなら YES を、できないなら NO を出力すること。

解法

出来上がる正方形の1辺の長さを $k$ とするとき、その正方形の面積は ${k}^{2}$ となります。

すべてのブロックを使用して正方形を作成することから、 \begin{align} {k}^{2} = {a}_{1} + {a}_{2} + \cdots + {a}_{n} \end{align} を満たすような整数 $k$ が存在するとき、正方形を作成できるということになります。

従って、与えられた数列 $a$ のすべての要素の和を求め、その値が平方数であるかどうかを判定することにより、この問題は解くことができます。

ただし、この問題では $n \leq 2 \times {10}^{5}$ であり、 ${a}_{i} \leq {10}^{9}$ であるため、数列 $a$ の要素の和は最大で $2 \times {10}^{14}$ となります。そのため、 $k$ について $1$ から $1.4 \cdot {10}^{7}$ 辺りまで愚直に順番に調べると、実行時間制限に間に合わないおそれがあります。

そこで、二分探索を用いてそのような $k$ の値が存在するかどうかを調べます。

具体的には、 ${k}^{2} \leq {a}_{1} + \cdots + {a}_{n} \lt {(k + 1)}^{2}$ が成立する整数 $k$ を求め、その $k$ に対して ${k}^{2} = {a}_{1} + \cdots + {a}_{n}$ が成立するかを調べることにより、問題の答えを導くことができます。

これは、

$l = 0, r = {10}^{9}$ として整数 $l, r$ を初期設定する
となるまで、以下の操作を繰り返す
- とするとき、
  - ${\mathrm{mid}}^{2} \leq {a}_{1} + \cdots + {a}_{n}$ が成立するならば、 $l \leftarrow \mathrm{mid}$ として更新する
  - そうでないならば、 $r \leftarrow \mathrm{mid}$ として更新する
最終的な $l, r$ が ${l}^{2} \leq {a}_{1} + \cdots + {a}_{n} \lt {r}^{2}$ となるため、 ${l}^{2} = {a}_{1} + \cdots + {a}_{n}$ が成立するかどうかを判定する

という手順によって、計算量を削減して判定を行うことができます。

ソースコード

正方形を作れるかどうかを判定し、作成できるならば true を、できないならば false を返す関数 solve() を以下のように実装しました。

int n;
ll a[int(2e5 + 5)];

bool solve() {
  ll sum = 0; // 数列 a の要素の合計値を表す整数 sum
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
    sum += a[i]; // sum に a[i] を足して更新する
  }

  // 二分探索を行う
  ll l = 0, r = 1e9;
  while (r - l > 1) {
    ll mid = (l + r) / 2;
    if (mid * mid <= sum) {
      l = mid;
    } else {
      r = mid;
    }
  }

  return l * l == sum; // 最終的に得られる l を用いて判定を行う
}