概略：\(\max A\) が小さいので、エラトステネスの篩のようにやれば \(O(N+\max A)\) でできます。

\(s_x\) を「\(x\) を平方数で割れるだけ割った残り」とします。例えば\(s_{12}=3\)、\(s_{72}=2\) です。

エラトステネスの篩のような以下のアルゴリズムにより、\(s_1,\ldots,s_X\) を \(O(X)\) 時間で列挙することができます。

for i = 1 to X:
  s[i]=i
for d = sqrt(X) to 2:
  for k = 1 to X/d/d:
    if s[k*d*d]%(d*d)==0:
      s[k*d*d]/=d*d

※ \(\frac{1}{1^2}+\frac{1}{2^2}+\frac{1}{3^2}+\ldots=O(1)\)

こうして計算した \(s\) を用いて公式解説と同様のことを行うことで、全体で \(O(N+\max A)\) 時間で答えを求めることができます。

実装例(C++)