Codeforces Round #842 (Div. 2) C. Elemental Decompress

Codeforces Round #842 (Div. 2) C. Elemental Decompress

Two cases produce no answers:

• One element appears more than twice in $a$
• After sorting, there is some index that $a[i]<i$ ($1$-indexed).

Consider there is some index that $a[i]<i$, then both $p[i]<i$ and $q[i]<i$ must satisfy. This is also true for the first $i-1$ index, so the numbers that are smaller than $i$ in both $p$ and $q$ are $(i-1)\times 2+2=i\ast 2$. This is a contradiction.
Otherwise, solutions always exist. One method is to constructively attach each element in $a$ to $p$ or $q$:

• Traverse from the biggest element to the smallest in $a$, if that number haven’t appeared in $p$ then attach it to $p$, otherwise attach it to $q$.
• Traverse from the biggest element to the smallest in $a$ again, if we attached it to $p$, find the biggest number that did not appear in $q$ and attach to $q$, vice versa.

A naive solution requires the $O(n^2)$ method to solve. We can reduce to $O(nlogn)$ by sorting elements in $a$ as pairs <element, index>.

Time complexity: $O(nlog(n))$

#include<bits/stdc++.h>
using namespace std;

int main()
{
    auto solve=[&](){
        int n;cin>>n;
        vector<int> a(n+1),st(n+1),p(n+1),q(n+1);
        for(int i=1;i<=n;i++)
        {
            cin>>a[i];
            st[a[i]]++;
        }
        if(st[1]>1)
        {
            cout<<"NO"<<endl;
            return ;
        }
        for(int i=1;i<=n;i++)
            if(st[i]>2)
            {
                cout<<"NO"<<endl;
                return ;
            }

        vector<int> pst(n+1),qst(n+1);
        for(int i=1;i<=n;i++)
        {
            if(pst[a[i]]&&qst[a[i]])
            {
                cout<<"NO"<<endl;
                return ;
            }
            if(pst[a[i]])
            {
                q[i]=a[i];
                qst[a[i]]=1;
            }
            else
            {
                p[i]=a[i];
                pst[a[i]]=1;
            }
        }

        multiset<int> m1,m2;
        for(int i=1;i<=n;i++)
        {
            if(!pst[i]) m1.insert(i);
            if(!qst[i]) m2.insert(i);
        }

        for(int i=1;i<=n;i++)
        {
            if(!p[i])
            {
                if(*m1.begin()>q[i])
                {
                    cout<<"NO"<<endl;
                    return ;
                }
                auto pos=m1.upper_bound(q[i]);
                p[i]=*(--pos);
                m1.erase(pos);
            }
            else
            {
                if(*m2.begin()>p[i])
                {
                    cout<<"NO"<<endl;
                    return ;
                }
                auto pos=m2.upper_bound(p[i]);
                q[i]=*(--pos);
                m2.erase(pos);
            }
        }

        cout<<"YES"<<endl;
        for(int i=1;i<=n;i++)
            cout<<p[i]<<" ";
        cout<<endl;
        for(int i=1;i<=n;i++)
            cout<<q[i]<<" ";
        cout<<endl;
    };

    int T;cin>>T;
    while(T--) solve();
    return 0;
}