第 5 回ドワンゴからの挑戦状予選 C

C - k-DMC

問題

長さ $n$ の英大文字からなる文字列が与えられる．これに対し，以下の形式で与えられる $q$ 個のクエリに答えよ．

各クエリでは，整数 $K$ が与えられる．このとき，以下を全て満たす整数の組 $(a, b, c)$ の数を求めよ．

$1 \leq a \lt b \lt c \leq n$
$s_a s_b s_c =$ DMC
$c - a \lt K$

制約

$3 \leq n \leq 10^6$
$1 \leq q \leq 75$
$3 \leq K \leq n$

考察

$D = \{ i \mid s_i =$ D $\}$ とする． $M, C$ も同様．

そして，以下の 3 つの累積和的な数列を考える．

$dsum_i = | \{ j \in D \mid 1 \leq j \leq i \} |$
$msum_i = | \{ j \in M \mid 1 \leq j \leq i \} |$
$dmsum_i = \sum \{ dsum_j \mid 1 \leq j \leq i, j \in M \} = | \{ (j, k) \in D \times M \mid 1 \leq k \leq j \leq i \}|$

これらを使うと， $s_i =$ Cのときにこれを末尾に持つような組の数は以下の式で求められる．

$$ (dmsum_i - dmsum_{i - K}) - (msum_i - msum_{i - K}) \cdot dsum_{i - K} $$

第 1 項で M が $(i-K, i]$ にあるような DM の個数を求めている．これでは D が $[0, i-k]$ にあるようなものも含まれるので，それを第 2 項で除いている．

後はこれを $i \in C$ について足し上げればいい．

実装例

提出 #8121212 - 第5回ドワンゴからの挑戦状予選

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45


#include <iostream>
#include <algorithm>
#include <vector>
#include <string>

using lint = long long;

int main() {
    int n;
    std::string s;
    std::cin >> n >> s;

    std::vector<lint> dsum(n + 1, 0), msum(n + 1, 0), dmsum(n + 1, 0);
    std::vector<int> cpos;
    for (int i = 1; i <= n; ++i) {
        char c = s[i - 1];
        dsum[i] = dsum[i - 1];
        msum[i] = msum[i - 1];
        dmsum[i] = dmsum[i - 1];

        if (c == 'D') {
            ++dsum[i];
        } else if (c == 'M') {
            ++msum[i];
            dmsum[i] += dsum[i];
        } else if (c == 'C') {
            cpos.push_back(i);
        }
    }

    int q;
    std::cin >> q;
    for (int p = 0; p < q; ++p) {
        int k;
        std::cin >> k;
        lint ans = 0;
        for (auto i : cpos) {
            int l = std::max(0, i - k);
            ans += (dmsum[i] - dmsum[l]) -
                   (msum[i] - msum[l]) * dsum[l];
        }
        std::cout << ans << std::endl;
    }
    return 0;
}