けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

Codeforces Round #489 (Div. 2) D. Nastya and a Game (R2100)

方針立ってからも、ちょっと実装辛い ^^;

問題へのリンク

問題概要

長さ  N の正の整数列  a_{1}, a_{2}, \dots, a_{N} と、正の整数  K が与えられる。

正数列の連続する部分列であって、その積が、和の  K 倍となっているものが何通りあるかを求めよ。

制約

  •  1 \le N \le 2 \times 10^{5}
  •  1 \le K \le 10^{5}
  •  1 \le a_{i} \le 10^{8}

考えたこと

パッと思うことは、積の中に 2 以上の整数ってあんまり多く含めることはできない。ちゃんと見積もると、積として考えられる最大値は

 2 \times 10^{5} \times 10^{5} \times 10^{8} = 2 \times 10^{18}

くらい。よって、2 以上の整数を 61 個くらい掛け算するとこれを超えてしまうので、そういうのを考える必要がない。

よって、数列を「1 以外の部分の index」のみを取り出してあげた場合には、幅が 61 以下くらいのみを考えればよいということで、十分間に合う。

コーナーケースとして、 K = 1 のときは、「単体の 1」が条件を満たすので別途数える。

#include <bits/stdc++.h>
#include <sys/time.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

int N;
long long K;
vector<long long> A;
using pll = pair<long long, int>;

long long solve() {
    long long asum = 0;
    for (int i = 0; i < N; ++i) asum += A[i];
    long long LIM = 2100000000000000000LL; //asum * K + 1;

    vector<pll> v;
    v.emplace_back(100, -1);
    for (int i = 0; i < N; ++i) if (A[i] > 1) v.emplace_back(A[i], i);
    v.emplace_back(100, N);
    long long res = 0;
    if (K == 1) res = N - ((int)v.size() - 2);

    for (int i = 1; i+1 < v.size(); ++i) {
        long long mul = 1;
        long long sum = 0;
        for (int j = i+1; j < v.size(); ++j) {
            if (mul > LIM / v[j-1].first) break;
            mul *= v[j-1].first;
            sum += v[j-1].first;
            if (mul % K != 0) continue;
            long long rsum = mul / K;
            long long insidesum = v[j-1].second - v[i].second - (j-1-i);
            long long need = rsum - sum - insidesum;
            if (need < 0) continue;
            long long left = v[i].second - v[i-1].second - 1;
            long long right = v[j].second - v[j-1].second - 1;
            if (left + right < need) continue;
            long long mi = max(0LL, need - right);
            long long ma = min(left, need);
            long long add = ma - mi + 1;
            res += add;
        }
    }
    return res;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0); 

    while (cin >> N >> K) {
        A.resize(N);
        for (int i = 0; i < N; ++i) cin >> A[i];
        cout << solve() << endl;
    }
}