a bit of codeforces: Round #268 (Div. 1) Problem B

Problem B:
468B - Two Sets

Summary:
Partition \(N\leq 10^5\) numbers into \(A\) and \(B\) such that if \(x \in A\) then \(a - x \in A\) and if \(x \in B\) then \(b-x \in B\) for two given constant \(a,b\).

Solution:
The numbers occur in pairs, and as such we can model the problem as a graph problem. Let's say that the partitioning is done by coloring the node with \(A\) or \(B\). We have two kinds of edge from each number \(x\) (as the vertex): an A-edge that connects \(a\) and \(x-a\), and/or a B-edge that connects \(b\) and \(x-b\). We can build this adjacency information in \(O(N\lg{N})\) time using binary search. Now to check whether we can color a node \(x\) with a color \(c\).
The color the nodes connected to \(x\) must have the same color as \(x\), since if it does not, then it can be quite easily shown that such coloring will lead to contradictory conclusions. Hence to complete the coloring, we have the following cases:
Case \(c = A\). (Case \(c=B\) is analogous)
1. node \(x\) does not have A-edge: then it is not possible to color the node with \(A\).
2. node \(x\) has an A-edge: then we continue this check recursively on the node \(a\), and the node \(b\) if it has one.
If the connected component satisfies this constraint, then we can color those nodes with color \(c\). Such checking will take \(O(N)\) time.

Implementation:

#include <iostream>
#include <cstdio>
#include <vector>
#include <algorithm>
#include <utility>
using namespace std;

int adj[100003][2];
int N, A, B;
vector<pair<int,int> > m;
vector<int> mm, pp;
int mark[100003];

bool check(int par, int u, int col){
    //printf("%d %d %d\n", par,u,col);
    mark[u] = col;
    int v = adj[u][0];
    int w = adj[u][1];
    bool ret = true;
    if(adj[u][col] == -1) return false;
    if(v != -1 && v != par) ret &= check(u, v, col);
    if(w != -1 && w != par) ret &= check(u, w, col);
    return ret;
}
    

int main(){
    cin >> N >> A >> B;
    m.clear();
    mm = vector<int>(N+3);
    pp = vector<int>(N+3);
    for(int i=0;i<N;++i){
        int u;
        cin >> u;
        m.push_back(make_pair(u, i));
    }
    sort(m.begin(), m.end());
    for(int i=0;i<N;++i){
        mm[i] = m[i].first;
        pp[m[i].second] = i;
        mark[i] = -1;
        adj[i][0] = adj[i][1] = -1;
    }
    bool ok = true;
    for(int i=0;i<N;++i){
        int lo = i, hi = N-1, mid;
        int u = A-mm[i];
        int v = B-mm[i];
        bool found = false;
        if(adj[i][0] != -1 || adj[i][1] != -1) found = true;
        while(lo <= hi){
            mid = (lo+hi)/2;
            if(mm[mid] < u){
                lo = mid+1;
            } else {
                hi = mid-1;
            }
        }
        if(mm[lo] == u){
            adj[i][0] = lo;
            adj[lo][0] = i;
            found = true;
        }
        lo = i, hi = N-1;
        while(lo <= hi){
            mid = (lo+hi)/2;
            if(mm[mid] < v){
                lo = mid+1;
            } else {
                hi = mid-1;
            }
        }
        if(mm[lo] == v){
            adj[i][1] = lo;
            adj[lo][1] = i;
            found = true;
        }
        if(!found){
            ok = false;
            break;
        }
    }
    /*
    for(int i=0;i<N;++i){
        printf("%d %d %d\n", i, adj[i][0], adj[i][1]);
    }
    */
    if(!ok) printf("NO\n");
    else {
        for(int i=0;i<N;++i){
            if(mark[i] != -1) continue;
            bool ch = check(-1, i, 0);
            if(ch) continue;
            ch = check(-1, i, 1);
            if(!ch){
                ok = false;
                break;
            }
        }
        if(!ok) printf("NO\n");
        else {
            printf("YES\n");
            for(int i=0;i<N;++i){
                if(i!=0) printf(" ");
                printf("%d", mark[pp[i]]);
            }
            printf("\n");
        }
    }
    return 0;
}