🕌

Kickstart 2020 RoundC: Perfect Subarray

2021/07/12に公開

問題

N個のinteger（negative valueの場合もあり）を含む配列が与えられたとき、perfect subarrayが幾つあるかを求める問題。
連続するsubarrayの合計値がperfect squareである時、perfect subarrayであると定義。perfect squareは、non negativeなintegerのそれ自身との積である数のことを言う。例えば、0, 1, 4, 9, 16など。

アルゴリズム

全てのsubarrayを求めて、各subarrayがperfectであることを確認すれば、答えを求めることはできます。subarrayは開始位置と終了位置で決まるので、開始位置が0,1,2,...,N-1のときに、終了位置はN-1個,N-2個,...,1個となり、それぞれのsubarrayについてsumを計算すると、全体で $O(N^3)$ の計算量となります（事前に累積和を求めておけば $O(N^2)$ で全subarrayをチェックできる）。
$1 \leq N \leq 10^5$ となるTest set2をpassするためには、別のアプローチが必要です。

与えられたarrayをAとして、Aの中の位置iについて累積和 $\rm{sum}(A[0...i])$ を求める。0からi-1までの位置kの $\rm{sum}(A[0...k])$ に対して、 $\rm{sum}(A[0...i])-perfect\ squareな値=\rm{sum}(A[0...k])$ となったとすると、kからiまでのsubarrayがperfect squareとなるはずです。
従って、Aの各位置iについて累積和を求めながら、上記の計算を行い、求めた値と等しい累積和となる位置kの個数をカウントに加えていく。累積和を求めていく時には、累積和をキー、同じ累積和を持つ位置kの数を値としてもつハッシュテーブルを作成して参照するようにし、またperfect squareを用いた上記の計算は、累積和の最小値より小さくなったら計算を打ち切って、次の位置へ移動する。

#include <bits/stdc++.h>
typedef long long ll;
using namespace std;


int main()
{
    ll T, N;
    cin >> T;
    for (ll t = 1; t <= T; ++t)
    {
        cin >> N;
        vector<ll> arr;
        for (ll i = 0; i < N; ++i)
        {
            int d;
            cin >> d;
            arr.push_back(d);
        }
        ll res = 0;
        ll sum = 0;
        ll min_sum = 0;
        unordered_map<ll, ll> m;
        m[0] = 1;
        for (ll i = 0; i < arr.size(); ++i)
        {
            sum += arr[i];
            min_sum = min(sum, min_sum);
            for (ll j = 0; sum - (j * j) >= min_sum; ++j)
            {
                if (m.count(sum - (j * j)) == 1)
                {
                    res += m[sum - (j * j)];
                }
            }
            m[sum] += 1;
        }
        cout
            << "Case #" << t << ": " << res << endl;
    }
}

計算量は、Aの各要素のループと、累積和を用いた計算のループで、後者のループでは累積和の最大値はN*MAX_Aで、またその時のsquareの個数は $\sqrt{N*MAX\_A}$ なので、 $O(N*\sqrt{N*MAX\_A}$ となる。また、累積和の計算や、最小の累積和の更新、ハッシュテーブルの参照はO(1)。

参考

こちらを参考にしました。
https://youtu.be/Lyay4WJwXyI

問題

アルゴリズム

参考

Discussion