問題概要

長さ $ N (\leq 2 \times 10^5)$ の整数列 $ A_1 = X, A_2 = X+D, \cdots, A_{N-1} = X+(N-2)D, A_{N} = X+(N-1)D $ がある (ただし $ X, D \leq 2 \times 10^8$ )．

この数列の要素を2つのグループ $ s, t $ に分けたとき，それぞれのグループに含まれる要素の和を $ S, T$ とする．$ S-T $ の値としてありうるものが何通りあるかを求めよ．

解法概要

$ D \lt 0 $ のときは $ X, D $ に -1 をかけるなどしても通り数が変わらないので，以下は $ D > 0 $ の場合を考える．D = 0 の場合は最後に補足する．

$ S+T = U $ とおくと，$ S-T = S - (U-S) = 2S - U $ となる．U が一定であることを考えると，この問題は $ S $ として取りうる値が何通りあるのかを考える問題と等価であることがわかる．

s の要素数を K とすると $ S = XK + D \sum_{A_i \in s} (i-1) $ という形で表現できる．$ \sum_{A_i \in s} (i-1) $ の値は最小値 $ L = \sum_{i=0}^{K-1} i $ から最大値 $ R = \sum_{i=0}^{K-1} (N-i) $ まで1ごとに変化させることができるので，s の要素数を K と固定したとき，$ S $ の和としてありうるのは $ R-L+1 $ 通りということになる．

ただし気にしなければならないのは，同じ和の値に対して s の要素の選び方が複数存在するような場合である． $ S = XK + D \sum_{A_i \in s} (i-1) $ の第 2 項についてはどうあがいても D の倍数にしかならないので，要素の選び方が違うのに同じ和になりうるのは $ S \% D = XK \% D $ の値が同じ場合である．

$ S \% D $ の値が s の要素数 K にしか拠らない (ただし入力で与えられる数は除いて) ことがわかったので，あとは $ S \% D = XK \% D $ の値ごとに $ S $ として取りうる値の区間を求めて，同じく$ S \% D $ の値ごとに数え上げていく．C++ では (負数) % (正数) が負 (か0) になるので，関数 modulo で補正している．商についても，正となるように補正した剰余と矛盾しないように関数 division で補正している．

複数の区間から取りうる値をカウントする方法はおそらくいくつか存在するが，私の実装した方針を説明する．

まず，区間を始点でソートしておく．注目する区間が今までの区間とまったく被らなければ，その区間の取りうる値の数を答えに足しこむ．注目する区間の始点側が今まで見た区間と被っていれば，今まで見た区間の最大値から注目する区間の最大値の範囲の値の数を答えに足しこむ．区間は始点でソートしているので，終点側でだけ今まで見た区間と被ることはない．最後に，注目する区間が今まで見た区間と丸かぶりしていれば，何も足さない．

これで $ D > 0 $ の場合は計算できるが，これまでの計算は D による除算を用いるため D = 0 のときは適用できない．

D = 0 の場合は，同じ値が N 個続く整数列が与えられたことになるので， S として取りうる値はその値を 0 個〜 N 個取った場合の $ N+1 $ 通りである．ただし， $ D = 0, X = 0 $ の場合はいくつの要素を取っても和が 0 なので答えは 1 通りになる．

ソースコード

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

using namespace std;

using ll = long long;
using Pll = pair<ll, ll>;

ll INF = 1LL << 60;

ll division(ll a, ll b) {
    if(a < 0) {
        return - ((-a + b - 1) / b);
    }
    return a / b;
}

ll modulo(ll a, ll b) {
    return ((a % b) + b) % b;
}

int main() {
    ll n, x, d;
    cin >> n >> x >> d;

    if(d == 0) {
        cout << ((x == 0) ? ll(1) : n+1) << endl;
        return 0;
    } else if(d < 0) {
        x *= -1;
        d *= -1;
    }

    map<ll, vector<Pll>> mp;
    for(ll i=0;i<=n;++i) {
        ll l = (i-1) * i / 2, r = (n-i) * i + l;
        mp[modulo(i*x, d)].emplace_back(division(i*x, d) + l, r - l + 1);
            // (区間の最小値を D で割った値, その区間に含まれる和の個数) のペアを格納
    }

    ll ans = 0LL;
    for(auto itr=mp.begin();itr!=mp.end();++itr) {
        vector<Pll> v(itr->second);
        sort(v.begin(), v.end());

        ll max_r = -INF;
        for(Pll p: v) {
            if(max_r < p.first) {
                ans += p.second;
            } else {
                ans += max(0LL, p.second - (max_r - p.first + 1));
            }
            max_r = max(max_r, p.first + p.second - 1);
        }
    }

    cout << ans << endl;
}

所感

負数の除算は罠……．区間をソートして以降の数え上げ方も地味に苦労した．

NoiminのNoise

AtCoder Beginner Contest 147 F - Sum Difference

問題概要

解法概要

ソースコード

所感