**Author:** Kamil Debowski

**Tester:** Praveen Dhinwa

**Editorialist:** Praveen Dhinwa

### PROBLEM

There are n cells in a row. Some of these cells are empty and the others are occupied by the soldiers. You are provided this information in an array a of length n, where a_i is binary number: a_i = 0 denotes that the i-th cell is empty, whereas a_i = 1 says that the i-th cell is occupied by a soldier.

You are playing a game in which you are allowed to make following kind of move.

Choose a soldier (it takes exactly one second to choose a soldier) and ask him to keep moving to the right till he can, i.e. he stops only if he is on the last cell or the next cell to him is already occupied by some soldier. Each soldier moves at a speed of one cell to the right per second. You are allowed to choose a soldier only if he is able to move at least one cell to his right.

You want to play the game as long as possible. Find out the maximum number of seconds for which you can enjoy the game.

### PREREQUISITES

greedy algorithms, understanding of proving the optimality of greedy algorithms, understanding of computing prefix sum of an array

### QUICK SOLUTION

We repeatedly apply the following strategy till we can:

- Choose the left most soldier that can be moved. We move it to its right as far as possible.

Please refer to next section for checking proof of correctness of this strategy.

The direct implementation of this strategy will take \mathcal{O}(n^2) time. We will need to make below observations to optimize it to \mathcal{O}(n).

For each move, we think of it as making two separate moves, one is to just choose the soldier, and the other will be to moving the solider to the cells to its right.

- Total number of cells moved by each soldier will remain constant regardless of the way you choose the soldiers to move.
- This means that the only part that is important for maximizing time taken for the game is to choose the soldiers as many as times as possible.

You can find the count of maximum number of times a solider can be chosen to move as follows. It will be equal to total number of consecutive groups of zeros that lie after the soldier (i.e. are in the part of the row that is right to the soldier). You can compute this information in \mathcal{O}(n) time. You can find the total number of cells moved by all the soldiers in \mathcal{O}(n) time too. Thus, overall it leads to \mathcal{O}(n) time.

### SOLUTION

Let us take examples and make some observations from them.

00111: In this case it is not possible to choose any soldier for making a move. So, the answer is zero. This tells us if all the soldiers have moved the to the right of the row, then we can’t make any move and answer will be zero.

01011: Here, we can only move the solider from 2^{nd} cell to the 3^{rd} cell. It takes one second to choose and one second to move 1 cell to the right. So overall it will take 2 seconds. Now, the situation becomes : 00111, and no further moves are possible.

1010: We have two soldiers that can be chosen to be moved in the first move. Let us analyze both the cases.

- Suppose we choose the soldier at first cell to move, we get 0110 (it takes 2 seconds). Now, we can only move the solider at 3^{rd} cell to the right, we get 0101 (2 seconds again). At last we move the soldier at second cell to the right and get 0011. Overall we took 2 + 2 + 2 = 6 seconds.
- Suppose we choose to move the soldier at 3rd cell instead. We get 1001 (it takes 2 seconds). Now, we move the soldier at first cell to 2 cells right (takes 3 seconds). The game lasted for 5 seconds.

This examples makes us to think that what was special about moving the soldier at first cell first than moving the soldier at the third cell. Why did choosing the soldier at first cell made the game go longer?

Let us think about move in a different way. Each move has two parts, choosing a soldier and moving the solider to some cells to the right. Let us analyze these steps separately. In the first case, we had 3 soldier choosing steps, and overall we moved soldiers 3 cells to right. In the second case, the soldier choosing steps were only 2, where as the total number of cells moved to the right by the soldiers were 3. You can see that total number of cells moved to the right by the soldiers are same in both the cases. In fact, if you ponder more over it, you realize that this fact is indeed true for any general case. We know that in the final configuration, all the soldiers will be on the right of the row. During the game the soldiers can’t change their relative positions, so you can uniquely identify the final positions of each soldier. So, the number of cells that a solider has to move will be the distance between his initial and final cells. As the initial and final positions of each soldier is unique, so the number of cells moved to the right will remain unique for a soldier regardless of order of the moves.

This means that in order to maximize the time taken in the game, we should choose the soldiers as many times as possible. Let us understand what’s the maximum number of times we can possibly choose a soldier.

Let us say that we have following example row: 0010\textbf{1}0011000101. For the soldier shown in bold, we want to figure out what is the maximum number of times we can choose this soldier. We can see that there are three group of consecutive zeros that occur after the soldier’s cell 00101\textbf{00}11\textbf{000}1\textbf{0}1 (shown in bold). We will prove that we can choose this soldier at most 3 times.

You can notice that it is impossible to increase the number of these groups of consecutive zeros by our moves. This is because if we choose a soldier to move, we have to move it to the right cell till he can move, so we will never leave the soldier in middle of zeros creating more number of consecutive groups.

Also, if you choose the soldier and decide to move him to the right, then you will be decreasing the number of these consecutive zeros by one (001010\textbf{1}11\textbf{000}1\textbf{0}1, 2 groups of zeros). The group of zeros that lies right next to the soldier will no longer be counted after the movement. Hence, you can only move the soldier only by the number of consecutive groups of zeros that are to the right of him.

So, this will be our strategy for achieving the maximum time in the game. We pick the leftmost soldier that can be moved and move him. We keep on doing this till we move all the soldiers to the right of row. This will be optimal as in this case the soldiers will be chosen maximum number for the moves. To prove it, we have to show that in this case, the total number of times a soldier is chosen will be equal to number of groups of consecutive zeros to the right of the soldier. The most important observation is to realize that for all the soldiers except the one that we are moving, there is no change in number of consecutive group of zeros after them. The number of groups for the soldier we are moving decreases by exactly one. This is why our solution is optimal.

For solving the second subtask, you can directly simulate this observation in \mathcal{O}(n^2) time. We need to do something better for the last subtask.

For each soldier, we find the number of consecutive groups of zeros that lie to the right of him. Let cnt_i denote this for count for i-th soldier. We go from right to the left, we will say that a group of zeros has started only if the last encountered cell was 1 and the current cell is zero.

```
cnt[n] = a[i] == 0 ? 1 : 0;
for i = n - 1 to 1:
cnt[i] = cnt[i + 1]
if (a[i] == 0 && a[i + 1] == 1):
cnt[i]++; // a new group of consecutive zeros has begun.
```

We will also need to final the total number of cells moved by the soldiers to find the total time. Let final_cell[i] denote the final cell at which the i-th soldier is supposed to end up.

```
final_cell[n] = n;
soldiers = 0;
for i = n to 1:
final_cell[i] = n - soldiers;
if (a[i] == 1):
soldiers++;
total = 0
for i = 1 to n:
if (a[i] == 1):
total += final_cell[i] - i;
```

Finally, our answer will be total number of choose moves and the total number of cells moved

```
ans = total;
for i = 1 to n:
if (a[i] == 1):
ans += cnt[i];
```

Time complexity of this approach is \mathcal{O}(n) per test case.