간식 시간 (JUNGOL 5721)

https://jungol.co.kr/problem/5721

---

아이디어

어떤 구간 [x+1,y]가 가능하기 위해서는 다음중 하나를 만족해야합니다. 이때 I[i]를 1부터 i까지 'I'의 개수, C[i]를 1부터 i까지의 'C'의 개수, P[i]를 1부터 i까지의 'P'의 개수라고 합시다.

1. I[y]-I[x]==C[y]-C[x]<P[y]-P[x]

2. P[y]-P[x]==C[y]-C[x]<I[y]-I[x]

3. I[y]-I[x]==P[y]-P[x]<C[y]-C[x]

 

1번 식을 정리하면 I[y]-C[y] 와 I[x]-C[x]가 같은 x값을 찾음과 동시에 P[y]-I[y]보다 P[x]-I[x]가 더 작아야합니다. 이러한 모든 x에 대하여 그 위치 x에서의 dp값을 더해주면 됩니다.

 

이는 세그먼트 트리로 해결해줄 수 있습니다.

I[y]-C[y], P[y]-I[y]와 같은 값은 미리 구해놓을 수 있으므로, 모든 {I-C,P-I], {P-I,C-P}, {C-P,I-C}의 값을 각각 벡터에 넣고 정렬한 뒤, 인덱스를 매겨줄 수 있습니다. 그러면 어떠한 y에 대해서 조건을 만족하는 x의 값을 O(log N)에 찾을 수 있고, 갱신까지 O(log N)에 가능해지게 됩니다.

---

소스코드

#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
#define ll long long
using namespace std;

ll tree[3][4040404]={0},mod=998244353;
void update(ll idx,ll s,ll e,ll node,ll loc,ll num)
{
    tree[idx][node]+=num;
    tree[idx][node]%=mod;
    if (s==e) return;

    ll mid=(s+e)/2;
    if (loc<=mid) update(idx,s,mid,node*2,loc,num);
    else update(idx,mid+1,e,node*2+1,loc,num);
}
ll Sum(ll idx,ll s,ll e,ll node,ll l,ll r)
{
    if (e<l || r<s) return 0;
    if (l<=s && e<=r) return tree[idx][node];

    ll mid=(s+e)/2;
    return (Sum(idx,s,mid,node*2,l,r)+Sum(idx,mid+1,e,node*2+1,l,r))%mod;
}

ll IC[1010101]={0},CP[1010101]={0},PI[1010101]={0};

int main()
{
    ios_base::sync_with_stdio(false); cin.tie(NULL);
    ll N,i;
    string s;
    vector <pair <ll,ll> > IisC,CisP,PisI;

    cin>>s;
    N=s.length(); s=" "+s;

    IisC.push_back({0,0});
    CisP.push_back({0,0});
    PisI.push_back({0,0});
    for (i=1;i<=N;i++)
    {
        IC[i]=IC[i-1];
        CP[i]=CP[i-1];
        PI[i]=PI[i-1];
        if (s[i]=='I') IC[i]++,PI[i]--;
        if (s[i]=='C') CP[i]++,IC[i]--;
        if (s[i]=='P') PI[i]++,CP[i]--;

        IisC.push_back({IC[i],CP[i]});
        CisP.push_back({CP[i],PI[i]});
        PisI.push_back({PI[i],IC[i]});
    }

    sort (IisC.begin(),IisC.end());
    sort (CisP.begin(),CisP.end());
    sort (PisI.begin(),PisI.end());

    IisC.erase(unique(IisC.begin(),IisC.end()),IisC.end());
    CisP.erase(unique(CisP.begin(),CisP.end()),CisP.end());
    PisI.erase(unique(PisI.begin(),PisI.end()),PisI.end());

    update(0,0,N,1,lower_bound(IisC.begin(),IisC.end(),pair <ll,ll> {0,0})-IisC.begin(),1);
    update(1,0,N,1,lower_bound(CisP.begin(),CisP.end(),pair <ll,ll> {0,0})-CisP.begin(),1);
    update(2,0,N,1,lower_bound(PisI.begin(),PisI.end(),pair <ll,ll> {0,0})-PisI.begin(),1);

    ll res=0,idx1,idx2,num;
    for (i=1;i<=N;i++)
    {
        res=0;

        //cout<<IC[i]<<' '<<CP[i]<<' '<<PI[i]<<"\n";

        idx1=upper_bound(IisC.begin(),IisC.end(),pair <ll,ll> {IC[i],CP[i]})-IisC.begin();
        idx2=upper_bound(IisC.begin(),IisC.end(),pair <ll,ll> {IC[i],2e9})-IisC.begin();
        num=Sum(0,0,N,1,idx1,idx2-1);
        res+=num;
        
        idx1=upper_bound(CisP.begin(),CisP.end(),pair <ll,ll> {CP[i],PI[i]})-CisP.begin();
        idx2=upper_bound(CisP.begin(),CisP.end(),pair <ll,ll> {CP[i],2e9})-CisP.begin();
        num=Sum(1,0,N,1,idx1,idx2-1);
        res+=num;
        
        idx1=upper_bound(PisI.begin(),PisI.end(),pair <ll,ll> {PI[i],IC[i]})-PisI.begin();
        idx2=upper_bound(PisI.begin(),PisI.end(),pair <ll,ll> {PI[i],2e9})-PisI.begin();
        num=Sum(2,0,N,1,idx1,idx2-1);
        res+=num;

        res%=mod;
        update(0,0,N,1,lower_bound(IisC.begin(),IisC.end(),pair <ll,ll> {IC[i],CP[i]})-IisC.begin(),res);
        update(1,0,N,1,lower_bound(CisP.begin(),CisP.end(),pair <ll,ll> {CP[i],PI[i]})-CisP.begin(),res);
        update(2,0,N,1,lower_bound(PisI.begin(),PisI.end(),pair <ll,ll> {PI[i],IC[i]})-PisI.begin(),res);

        //cout<<res<<"??\n";
    }

    cout<<res;
}