---
# System prepended metadata

title: Candy Rush

---

# [Candy Rush](https://po2punkt0.kattis.com/problems/candyrush)
We want to count the number of intervals such that if
* x is the number of caramels in the interval,
* y is the number of licorice in the interval,
* K is the cap on caramels in the interval,
then $\min(x, K) \geq \frac{x + y}{2}$.


## Lets do some case analysis depending on if $x$ goes above $K$
Case 1: $x \geq K$
Then $\min(x, K) = K$, so $$2 * K \geq x + y \implies x \geq y.$$

Case 2: $x < K$
Then $\min(x, K) = x$, so $$2 * x \geq x + y \implies x \geq y.$$

Notice that in both cases the condition $x \geq y$ has to be true. However in the first case this condition isn't sufficient so we are overcounting the intervals where:
* $x \geq y$, and
* $x + y > 2 * K$.
The condition $x \geq K$ is implied from the other two so it's reduntant. Also note that $x + y$ is simpily the length of the interval.


## Counting intervals where $x \geq y$
Define the array $$p[i] = x - y, \text{in the interval [0, i]}.$$
For any interval $[l, r]$: $$x - y = p[r] - p[l - 1].$$
Now to counting: $$x \geq y \iff p[r] \geq p[l - 1].$$
So if we iterate $r = 0 \dots n-1$, for each $p[r]$ we count the number of earlier $p$-values that are less than or equal to $p[r]$. This can be done efficiently using a Fenwick tree or segment tree.

## Counting the overcounted intervals
We will do the exact same but only insert elements when we are $2K$ positions ahead of them

# Implementation
:::spoiler Code 
```cpp=
#include <bits/stdc++.h>
using namespace std;

struct T {
    int n;
    vector<int> e;
    T(int n) : n(n), e(n, 0) {}

    void add(int i, int v) {
        for (; i < n; i |= i+1)
            e[i] += v;
    }

    int sum(int r) {
        int s = 0;
        for (; r >= 0; r = (r & (r + 1)) - 1)
            s += e[r];
        return s;
    }
};

int main() {
    cin.tie(0)->sync_with_stdio(0);

    int n, k;
    cin >> n >> k;

    vector<int> p(n + 1);

    string s;
    cin >> s;
    for (int i = 0; i < n; i++)
        p[i + 1] = p[i] + (s[i] == 'C' ? 1 : -1);

    long long ans = 0;

    T t(2 * n);
    t.add(n, 1);
    for (int i = 0; i < n; i++) {
        ans += t.sum(p[i + 1] + n);
        t.add(p[i + 1] + n, 1);
    }

    t = T(2 * n);
    for (int i = 0; i < n; i++) {
        ans -= t.sum(p[i + 1] + n);
        if (i + 1 >= 2 * k)
            t.add(p[i + 1 - 2 * k] + n, 1);
    }

    cout << ans << '\n';
}
```
:::
