**PROBLEM LINK:**

Practice

Contest

**Author:** Pankaj Jindal, Piyush Kumar

**Tester: **Sergey Kulik

**Editorialist:** Aman jain

# DIFFICULTY:

Medium

# PREREQUISITES:

Set, Segment Tree

# PROBLEM:

Given a array $A$ of size N each filled with a color with value $A_i$. There are $Q$ queries, each query update color of box at position $Pos_i$ with color $Col_i$. For each query you need to find number of ways to select $W$ adjacent boxes of same color from total of $N$ boxes.

# QUICK EXPLANATION:

Maintain a set of starting indices of color segments. Now on each query there are three possibilities, either we assign the same color, so no change to the array, or we have an update that breaks a continuous segment of colors, or an update that joins two continuous segments, each of which can be done in $O(log(N))$. Update your answer for the query and correspondingly update set of starting indices of color segments.

# EXPLANATION:

Let's start with easy constraints. Let say you have a color segment of length $L$, then number of ways to select $W$ adjacent boxes from this segment would be $max(0,(L-W+1))$. Now if you have sequence of segments say, $s1s2s3s4$, then number of ways to select $W$ boxes would be sum of number of ways for each segment i.e. $\sum{max(0,L_i-W+1)}$. Complexity of this solution would be $O(N*Q)$, this works within time-limit for first sub-task but will fail on main task.

To solve hard part efficiently, we can store the result of the query for the input string in $ans$ and maintain set that store starting indices of color segments. After each query operation we need to insert or erase indices from set accordingly.

If you analyse closely you will find that each update operation looks like:

a) if update operation, does not change the state of boxes from current state, then there will be no change in answer to this query.

b) state change can be described as, $......s_a c1 s_b......$ to $......s_a c2 s_b......$, where $s_a$ is segment to the left and $s_b$ segment to the right of update position $pos_i$, with box color changed from $c1$ to $c2$.

Following cases possible (assuming 0 based indexing):

1) $pos_i$ is 0 and $c1$ lies in $s_b$ i.e. same in color to $s_b$, after the operation splitting takes place and $ans = ans - ways(len(s_b)+1) + ways(len(1)) + ways(len(s_b))$, here ways(l) means number of ways to get $W$ adjacent boxes from segment of length $l$.

2) $pos_i$ is 0 and $c1$ does not lies in $s_b$, after updation $c2$ lies in $s_b$, so merging takes place. $ans = ans - ways(1) - ways(len(s_b)) + ways(len(s_b)+1)$.

3) $pos_i$ is 0 and $c1$ and $c2$ both does not lie in $s_b$, no change in answer.

4) $pos_i$ is $N-1$ and $c1$ lies in $s_a$, after the operation splitting takes place and $ans = ans - ways(len(s_a)+1) + ways(len(1)) + ways(len(s_a)).$

5) $pos_i$ is $N-1$ and $c1$ does not lies in $s_a$, after updation $c2$ lies in $s_a$, so merging takes place. $ans = ans - ways(1) - ways(len(s_a)) + ways(len(s_a)+1).$

6) $pos_i$ is $N-1$ and $c1$ and $c2$ both does not lie in $s_a$, no change in answer.

7) for rest of the cases $pos_i$ lies in between 0 and $N-1$, $s_a = s_b$, i.e. same in color and $c1$ lies in $s_a$, now splitting takes place, $ans = ans - ways(len(s_a)+1+len(s_b)) + ways(len(s_a)) + ways(1) + ways(len(s_b))$.

8) $s_a = s_b$, $c1$ does not lies in $s_a$, $c2$ lies in $s_a$, $ans = ans - + ways(len(s_a)) - ways(1) - ways(len(s_b)) + ways(len(s_a)+1+len(s_b)).$

9) $s_a = s_b$, $c1$ and $c2$ does not lies in $s_a$, no change in answer.

10) $s_a != s_b$, $c1$ lies in $s_a$ and $c2$ lies in $s_b$, $ans = ans - ways(len(s_a)+1)+ways(len(s_a))-ways(len(s_b))+ways(len(s_b)+1)$.

11) $s_a != s_b$, $c1$ lies in $s_a$ and $c2$ does not lie in both, $ans = ans - ways(len(s_a)+1) + ways(1) + ways((len(s_a))$.

12) $s_a != s_b$, $c1$ lies in $s_b$ and $c2$ lies in $s_a$, $ans = ans - ways(len(s_b)+1)+ways(len(s_b))-ways(len(s_a))+ways(len(s_a)+1)$.

13) $s_a != s_b$, $c1$ lies in $s_b$ and $c2$ does not lie in both, $ans = ans - ways(len(s_b)+1) + ways(1) + ways((len(s_b))$.

14) $s_a != s_b$, $c1$ does not lies in both $s_a$ and $s_b$, $c2$ lies in $s_a$, $ans = ans - ways(1)-ways(len(s_a))+ways(len(s_a)+1)$.

15) $s_a != s_b$, $c1$ does not lies in both $s_a$ and $s_b$, $c2$ lies in $s_b$, $ans = ans - ways(1)-ways(len(s_b))+ways(len(s_b)+1)$.

16) $s_a != s_b$, $c1$ and $c2$ does not lies in both $s_a$ and $s_b$, no change in answer.

# COMPLEXITY:

Each query operation requires insertion, deletion of elements from set and answer update, all can be done in $O(log(N))$. For $Q$ queries, overall complexity would be $O(Q*log(N))$.

# ALITER:

Tester Sergey also suggested that above problem can be solved using Segment Tree, which would be much modular approach than handling many cases. **HINT:** For each node of segment tree stores length of color segment from left and length of color segment from right, color of left and right segment and ans for query for segment range. Update these values accordingly on each query operation.

If you are not able to solve the problem using Segment Tree, go through Sergey code in refernces.

**AUTHOR'S, TESTER'S and Editorialist's SOLUTIONS:**

author

tester

editorialist