AGC023-A Zero-Sum Ranges
(Test/AtCoder/agc023_a.test.cpp)

• View this file on GitHub
• Last update: 2023-07-22 13:55:49+09:00
• Problem: https://atcoder.jp/contests/agc023/tasks/agc023_a

「空でない連続する $A$ の部分列」を「非空な区間」と言い換える。 非空な区間 $[l, r)$ であって $\displaystyle \sum_{i = l}^{r - 1} A_i = 0$ を満たすものの数を数える問題となる。

これは、 $A$ の累積和 $S$ をとることで $S_{r} - S_{l} = 0$ となる区間の数を数える問題に言い換えることができる。式変形して $S_{r} = S_{l}$ と解釈する。

ここで、 $S_{i} = x$ を満たす $i$ の数が $k$ 個存在する時、 $S_{l} = x, S_{r} = x$ となるように $l, r$ を選択する通り数は $\binom{k}{2} = \frac{k \times (k - 1)}{2}$ である。

よって連想配列等を利用して $S_{i}$ を値の種類と数にまとめ上げることでこの問題を解くことができた。

この問題を部分問題に含む問題が現在まで沢山出題されているので、解法を残した。

Depends on

• 加法群 (Src/Algebra/Group/AdditiveGroup.hpp)
• 1次元累積和 (Src/DataStructure/PrefixSum1D/PrefixSum1D.hpp)
• 静的な列上の区間和クエリ (Src/DataStructure/PrefixSum1D/StaticRangeSumSolver.hpp)
• 標準データ型のエイリアス (Src/Template/TypeAlias.hpp)

Code

#define PROBLEM "https://atcoder.jp/contests/agc023/tasks/agc023_a"

#include "../../Src/Template/TypeAlias.hpp"
#include "../../Src/DataStructure/PrefixSum1D/StaticRangeSumSolver.hpp"

#include <iostream>
#include <vector>
#include <map>

using namespace zawa;

i32 main() {
    std::cin.tie(nullptr)->sync_with_stdio(false);

    u32 N; std::cin >> N;
    std::vector<i64> A(N);
    for (auto& a : A) std::cin >> a;

    Ruisekiwa<i64> S(A);
    std::map<i64, i32> mp{};
    u64 ans = 0;
    for (const auto v : S) {
        ans += mp[v];
        mp[v]++;
    } 

    std::cout << ans << std::endl;
}

#line 1 "Test/AtCoder/agc023_a.test.cpp"
#define PROBLEM "https://atcoder.jp/contests/agc023/tasks/agc023_a"

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 2 "Src/DataStructure/PrefixSum1D/StaticRangeSumSolver.hpp"

#line 2 "Src/Algebra/Group/AdditiveGroup.hpp"

namespace zawa {

template <class T>
class AdditiveGroup {
public:
    using Element = T;
    static constexpr T identity() noexcept {
        return T{};
    }
    static constexpr T operation(const T& l, const T& r) noexcept {
        return l + r;
    }
    static constexpr T inverse(const T& v) noexcept {
        return -v;
    }
};

} // namespace zawa
#line 2 "Src/DataStructure/PrefixSum1D/PrefixSum1D.hpp"

#line 4 "Src/DataStructure/PrefixSum1D/PrefixSum1D.hpp"

#include <cmath>
#include <vector>
#include <cassert>
#include <algorithm>
#include <type_traits>
#include <functional>

namespace zawa {

template <class Group>
class PrefixSum1D {
private:
    using T = typename Group::Element;
    std::vector<T> dat_;

    constexpr bool rangeCheck(u32 l, u32 r) const {
        return (l <= r and r < dat_.size());
    }

public:
    PrefixSum1D() = default; 
    PrefixSum1D(const std::vector<T>& A) : dat_(A.size() + 1, Group::identity()) {
        dat_.shrink_to_fit();
        for (u32 i = 0 ; i < A.size() ; i++) {
            dat_[i + 1] = Group::operation(dat_[i], A[i]);
        }
    }

    inline T operator[](u32 i) const {
        assert(i < dat_.size());
        return dat_[i];
    }

    inline usize size() const {
        return dat_.size();
    }

    T product(u32 l, u32 r) const {
        assert(rangeCheck(l, r));
        return Group::operation(Group::inverse(dat_[l]), dat_[r]);
    }

    u32 lowerBound(u32 l, u32 r, const T& v) const {
        assert(rangeCheck(l, r));
        T value = Group::operation(v, dat_[l]);
        return std::lower_bound(dat_.begin() + l, dat_.begin() + r, value) - dat_.begin();
    }

    u32 upperBound(u32 l, u32 r, const T& v) const {
        assert(rangeCheck(l, r));
        T value = Group::operation(v, dat_[l]);
        return std::upper_bound(dat_.begin() + l, dat_.begin() + r, value) - dat_.begin();
    }

    template <class F>
    u32 maxRight(u32 l, const F& f) const {
        static_assert(std::is_convertible_v<decltype(f), std::function<bool(T)>>, "f must be function bool(T)");
        assert(l < dat_.size());
        assert(f(Group::identity()));
        auto f_ = [&](const T& v) -> bool {
            return f(Group::operation(v, Group::inverse(dat_[l])));
        };
        return std::partition_point(dat_.begin() + l, dat_.end(), f_) - dat_.begin();
    }

    template <class F>
    u32 minLeft(u32 r, const F& f) const {
        static_assert(std::is_convertible_v<decltype(f), std::function<bool(T)>>, "f must be function bool(T)");
        assert(r < dat_.size());
        assert(f(Group::identity()));
        auto f_ = [&](const T& v) -> bool {
            return f(Group::operation(Group::inverse(v), dat_[r]));
        };
        return dat_.rend() - std::partition_point(dat_.rbegin() + (dat_.size() - r - 1), dat_.rend(), f_) - 1;
    }

    const auto begin() const {
        return dat_.begin();
    }

    const auto end() const {
        return dat_.end();
    }
};

} // namespace zawa
#line 5 "Src/DataStructure/PrefixSum1D/StaticRangeSumSolver.hpp"

namespace zawa {

    template <class T>
    using StaticRangeSumSolver = PrefixSum1D<AdditiveGroup<T>>;

    template <class T>
    using Ruisekiwa = PrefixSum1D<AdditiveGroup<T>>;

};
#line 5 "Test/AtCoder/agc023_a.test.cpp"

#include <iostream>
#line 8 "Test/AtCoder/agc023_a.test.cpp"
#include <map>

using namespace zawa;

i32 main() {
    std::cin.tie(nullptr)->sync_with_stdio(false);

    u32 N; std::cin >> N;
    std::vector<i64> A(N);
    for (auto& a : A) std::cin >> a;

    Ruisekiwa<i64> S(A);
    std::map<i64, i32> mp{};
    u64 ans = 0;
    for (const auto v : S) {
        ans += mp[v];
        mp[v]++;
    } 

    std::cout << ans << std::endl;
}

Back to top page