(POI) XVIII OI Round I Task Lollipop

https://szkopul.edu.pl/problemset/problem/YPme8cPuC1zbS3oA0euLxywx/site/?key=statement

magical linear time answer!

key point: the value can be only 1 or 2 !!!

#include <iostream>
#include <cstdio>
#include <cassert>
#include <cstring>
using namespace std;

const int N = 2000000 + 6;
const int M = 2000000 + 6;
const int INF = (1<<28);

int a[N];
char s[N];

int pre[N],suf[N];
int pos_suf[N];
int next_1[N];

int main ()
{
    int n,q;
    scanf("%d %d\n",&n,&q);
    scanf("%s",s);
    int tot = 0;
    for (int i=1;n>=i;i++)
    {
        a[i] = (s[i-1] == 'T'?2:1);
        tot += a[i];
    }
    int mn_1 = n+1;
    for (int i=1;n>=i;i++)
    {
        pre[i] = pre[i-1] + a[i];
        if (a[i] == 1)
        {
            if (mn_1 == n+1)
            {
                mn_1 = i;
            }
        }
    }
    memset(pos_suf,-1,sizeof(pos_suf));
    next_1[n+1] = n+1;
    for (int i=n;i>=1;i--)
    {
        suf[i] = suf[i+1] + a[i];
        pos_suf[ suf[i] ] = i;
        if (a[i] == 1)
        {
            next_1[i] = i;
        }
        else
        {
            next_1[i] = next_1[i+1];
        }
    }
    while (q--)
    {
        int k;
        scanf("%d",&k);
        if (k > tot)
        {
            puts("NIE");
        }
        else if (k == tot)
        {
            printf("%d %d\n",1,n);
        }
        else if (k == tot-1)
        {
            if (a[1] == 1)
            {
                printf("%d %d\n",2,n);
            }
            else if (a[n] == 1)
            {
                printf("%d %d\n",1,n-1);
            }
            else
            {
                puts("NIE");
            }
        }
        else if (k == tot-2)
        {
            if (a[1] == 2)
            {
                printf("%d %d\n",2,n);
            }
            else if (a[n] == 2)
            {
                printf("%d %d\n",1,n-1);
            }
            else if (a[1] == 1 && a[n] == 1)
            {
                printf("%d %d\n",2,n-1);
            }
            else
            {
                assert(0);
            }
        }
        else
        {
            int need_minus = tot - k;
            int pos = pos_suf[need_minus];
            if (pos != -1)
            {
                printf("%d %d\n",1,pos-1);
                continue;
            }
            need_minus--;
            pos = pos_suf[need_minus];
            if (mn_1 == 1)
            {
                printf("%d %d\n",2,pos-1);
                continue;
            }
            int right_2 = next_1[pos] - pos;
            int left_2 = mn_1 - 1;
            if (left_2 > right_2)
            {
                int delta = left_2 - right_2;
                if (next_1[pos] == n+1)
                {
                    puts("NIE");
                    continue;
                }
                else
                {
                    //consider need_minus = 11, 222221........2222211
                    printf("%d %d\n",right_2+2,next_1[pos]);
                    continue;
                }
            }
            else
            {
                //consider need_minus = 13, 221.....22222211
                printf("%d %d\n",mn_1 + 1,pos + left_2-1);
                continue;
            }
        }
    }
}