Реши на C++ за O(n): У Little X есть n различных ц...

Анализ задачи

Даны $n$ различных целых чисел $p_1, p_2, \dots, p_n$. Нужно разбить их все на два подмножества $A$ и $B$ так, чтобы выполнялись условия:

1. Если $x \in A$, то $a - x$ (если оно вообще есть среди данных $n$ чисел) тоже должно попасть в $A$.
2. Если $x \in B$, то $b - x$ (если оно есть среди данных $n$ чисел) тоже должно попасть в $B$.

Причём каждое из данных чисел должно быть либо в $A$, либо в $B$ (никаких «неиспользованных» не остаётся).

Однако из формулировки в задаче часто пропускается критическая деталь:

«Если $x \in A$, то $a-x$ должно принадлежать множеству $A$»
означает не только «если $a-x$ уже среди наших $n$ чисел, то оно тоже обязано быть в $A$», но и косвенно накладывает запрет класть $x$ в $A$, если $a-x$ не встречается в списке из $n$ чисел. Ведь мы можем раскладывать только данные $n$ чисел по множествам, а «требовать» в $A$ число, которого в исходном наборе нет, невозможно.

Иными словами:

• Число $x$ можно поместить в $A$ только если $a-x$ тоже содержится среди $p_1,\dots,p_n$. И при этом тогда автоматически $a-x$ тоже идёт в $A$.
• Число $x$ можно поместить в $B$ только если $b-x$ тоже содержится среди $p_1,\dots,p_n$. И при этом тогда автоматически $b-x$ тоже в $B$.

Если для какого-то $x$ оба варианта (положить в $A$ или в $B$) невозможны, значит задачу решить нельзя. Либо же при попытках удовлетворить все «цепочки» условий возникнет противоречие.

Примеры из условия

text
n=4, a=5, b=9
p = [2, 3, 4, 5]

• Посмотрим, кто может попасть в $A$:

  • $2\in A$ возможно, только если $5 - 2 = 3$ тоже есть среди чисел. Да, $3$ в списке есть, значит «теоретически» $2$ может пойти в $A$, и тогда $3$ обязана быть там же.
  • $3\in A$ возможно, только если $5 - 3 = 2$ тоже в списке. Да, значит $2$ и $3$ либо вместе в $A$, либо нет.
  • $4\in A$ возможно, только если $5 - 4 = 1$ есть среди наших чисел. Но $1$ нет в списке $\{2,3,4,5\}$. Значит нельзя класть $4$ в $A$.
  • $5\in A$ возможно, только если $5 - 5 = 0$ есть среди чисел. Но $0$ нет в списке. Значит нельзя класть $5$ в $A$.

• Посмотрим, кто может попасть в $B$:

  • $2\in B$ возможно, только если $9 - 2 = 7$ в списке. Но $7$ нет, значит $2$ в $B$ положить нельзя.
  • $3\in B$ возможно, только если $9 - 3 = 6$ в списке. $6$ нет, значит тоже нельзя.
  • $4\in B$ возможно, только если $9 - 4 = 5$ есть среди чисел. Да, $5$ есть. Значит $4$ может быть в $B$, тогда $5$ тоже обязательно в $B$.
  • $5\in B$ возможно, только если $9 - 5 = 4$ в списке. Да, $4$ есть, значит $4$ и $5$ либо вместе в $B$, либо нет.

Получается, что:

• $\{2,3\}$ должны попасть в $A$ (иначе их вообще «девать некуда»).
• $\{4,5\}$ должны попасть в $B$.

Это действительно согласовано и даёт разбиение:

В ответе выводят:

text
YES
0 0 1 1

(если нумерация p идёт как [2,3,4,5], то 0 означает «в A», 1 — «в B»).

text
n=3, a=3, b=4
p = [1, 2, 4]

• Для $1$:

  • Можно ли в $A$? Нужно, чтобы $3 - 1 = 2$ был в списке. Он есть, так что теоретически $1$ может пойти в $A$ при условии, что вместе с ним и $2$ тоже попадёт в $A$.
  • Можно ли в $B$? Нужно, чтобы $4 - 1 = 3$ был в списке. Но $3$ нет, значит в $B$ положить $1$ нельзя.

  Следовательно, $1$ обязан оказаться в $A$.

• Для $2$:

  • Может ли быть в $A$? Тогда надо, чтобы $3 - 2 = 1$ был в списке — да, $1$ есть, более того, мы уже «поняли», что $1$ уходит в $A$. Значит $2$ может тоже быть в $A$ (и тогда это согласуется с тем, что $1$ в $A$).
  • Может ли быть в $B$? Тогда надо, чтобы $4 - 2 = 2$ был в списке (он и есть, это «само на себя»). Формально это условие не противоречит: $2$ в $B$ «требует» $2$ в $B$. Это замкнуто, и само по себе не мешает.

  На первый взгляд $2$ мог бы пойти куда угодно, но поскольку $1$ уже обязан в $A$, а «$1$ в $A$ $\implies$ $3-1=2$ тоже в $A$», то это заставляет $2$ тоже быть в $A$.

• Для $4$:

  • Может ли быть в $A$? Нужно, чтобы $3 - 4 = -1$ был в списке, а $-1$ явно нет. Значит $4$ невозможно положить в $A$.
  • Может ли быть в $B$? Нужно, чтобы $4 - 4 = 0$ был в списке, а $0$ нет. Значит $4$ невозможно положить и в $B$.

Итого, $4$ не получается положить ни в $A$, ни в $B$. Поэтому задача не имеет решения, ответ «NO».

Итоговая проверка «можно ли распределить»

1. Для каждого числа $x$ из входного набора проверяем, какие варианты (0 = «A», 1 = «B») вообще возможны.

   • «$x$ может быть в $A$» тогда и только тогда, когда $(a-x)$ тоже присутствует среди наших $n$ чисел.
   • «$x$ может быть в $B$» тогда и только тогда, когда $(b-x)$ присутствует среди наших $n$ чисел.
     Если для какого-то $x$ оба варианта отпали, сразу «NO».

2. Далее нужно учесть, что если (по ходу решения) вдруг «$x$ точно пошёл в $A$», это автоматически заставляет $(a-x)$ тоже быть в $A$, и так далее каскадом. Аналогично для $B$. Если эти каскадные требования вступают в противоречие, то решения нет.

3. Если для каких-то чисел выбор не жёстко фиксирован (им можно и в $A$, и в $B$), мы можем выбрать любой из вариантов, но с учётом связей «в паре» (при выборе $x\in A$ обязательно «$a-x$ в $A$» и т.п.). Если каскад не даёт противоречий, то решение существует.

Для реализации за $O(n)$ удобно воспользоваться следующей схемой:

• Сначала подготовим:

  • Сохраним все $p_i$ в массив p.
  • Построим хеш-таблицу/словарь indexOf, чтобы быстро по значению находить индекс в p.
  • Заведём массив possible[i] (например, в виде битовой маски), который говорит, куда может пойти $p_i$. Изначально possible[i] = 0b11 (т.е. оба варианта разрешены).
  • Если $(a - p_i)$ не в наборе, то убираем «бит» 0b01 (то есть «в $A$» нельзя). Если $(b - p_i)$ не в наборе, убираем «бит» 0b10 (то есть «в $B$» нельзя). Если после этого possible[i] == 0, немедленно выводим «NO» и завершаемся.

• Далее фиксируем «жёстко определившиеся» элементы, для которых осталась ровно одна возможность (possible[i] == 0b01 или possible[i] == 0b10). Складируем их в очередь и обрабатываем по принципу «широкого поиска»:

  • Извлекаем элемент $x$ с его уже определённой меткой (допустим, $x$ пошёл в $A$). Тогда необходимо, чтобы $(a-x)$ тоже пошёл в $A$. Это значит, что у $(a-x)$ мы должны сохранять только бит «можно в $A$» и выкинуть «можно в $B$». Если у него это приводит к «нельзя никуда» — противоречие, «NO». Если у него раньше было «оба можно», а теперь стало «только $A$», то он тоже стал «жёстко определённым», добавляем его в очередь. Аналогично для $x\in B$.

• По ходу этого процесса либо выйдет конфликт, либо мы «сжав» все связи получим подмножество элементов, у которых осталось по одному варианту. Те элементы, у которых и в конце остались оба варианта $A$ или $B$, мы можем свободно (без противоречий) выбрать как $A$ или $B$. Обычно достаточно «пробежаться» и всем им присвоить, скажем, $A$ (если оно у них действительно в possible[i]), сделав аналогичную финальную «протяжку» по связям. В этой задаче не нужно искать все способы, достаточно любого корректного.

• Если в итоге все $p_i$ получили допустимые цвета без конфликта, печатаем «YES» и раскладку (0/1). Если по пути — «NO», выводим «NO».

Этот метод укладывается в $O(n)$–$O(n \log n)$ (все связи — максимум по 2 «партнёра» на каждый $x$: $a-x$ и $b-x$).

Решение на C++

Ниже приведён вариант кода, реализующий описанную идею.

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

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n;
    long long a, b;
    cin >> n >> a >> b;
    
    vector<long long> p(n);
    for(int i=0; i<n; i++){
        cin >> p[i];
    }
    
    // Сохраним в словарь "значение -> индекс"
    // Чтобы быстро находить, где лежит a-x или b-x
    unordered_map<long long,int> indexOf;
    indexOf.reserve(n);
    indexOf.max_load_factor(0.7f);
    for(int i=0; i<n; i++){
        indexOf[p[i]] = i;
    }
    
    // possible[i] будет битовой маской: 
    //    0b01 (единица) - можно ли i-й элемент положить в A,
    //    0b10 (двойка)  - можно ли i-й элемент положить в B.
    // Изначально ставим 0b11 (оба варианта), 
    // потом будем вычеркивать невозможное.
    vector<int> possible(n, 0b11);  
    
    // Предварительно обрежем те варианты, где a-x или b-x не существует в наборе
    // Если a-x не найдём, убираем бит "A"
    // Если b-x не найдём, убираем бит "B"
    for(int i=0; i<n; i++){
        long long x = p[i];
        
        // Проверка на возможность "x в A"
        // Нужно, чтобы (a - x) был в словаре
        if(!indexOf.count(a - x)){
            // Невозможно x в A
            possible[i] &= ~0b01; // вычёркиваем бит A
        }
        
        // Проверка на возможность "x в B"
        // Нужно, чтобы (b - x) был в словаре
        if(!indexOf.count(b - x)){
            // Невозможно x в B
            possible[i] &= ~0b10; // вычёркиваем бит B
        }
        
        // Если оказалось, что вариантов нет, сразу NO
        if(possible[i] == 0){
            cout << "NO\n";
            return 0;
        }
    }
    
    // Очередь для "жёстко зафиксированных" элементов
    // Будем хранить пары (индекс, color), где color=0 для A, color=1 для B.
    queue<pair<int,int>> q;
    
    // Если у кого-то один бит остался, сразу добавляем в очередь
    auto pushForced = [&](int i){
        // Если possible[i] == 1 (бит 0b01) => color=0
        // Если possible[i] == 2 (бит 0b10) => color=1
        // Если 3 - оба варианта, пока не трогаем
        if(possible[i] == 1){
            q.push({i,0});
        } 
        else if(possible[i] == 2){
            q.push({i,1});
        }
    };
    
    for(int i=0; i<n; i++){
        if(possible[i] == 1 || possible[i] == 2){
            pushForced(i);
        }
    }
    
    // Функция, которая "заставляет" элемент i идти в цвет c (0 или 1),
    // и сразу обрезает вариативность. Если конфликт - вернёт false.
    auto fixColor = [&](int i, int c){
        // c=0 => бит A=0b01, c=1 => бит B=0b10
        int bit = (1 << c); 
        // Если у i уже не включает bit, значит конфликт
        if((possible[i] & bit) == 0){
            return false;
        }
        // Иначе "обрезаем" всё, кроме нужного бита
        if(possible[i] != bit){
            possible[i] = bit;
            // Добавим в очередь, чтобы отработать связанные
            q.push({i, c});
        }
        return true;
    };
    
    // Выполним обход в ширину, отрабатывая "жёсткие" приписки
    while(!q.empty()){
        auto [i, c] = q.front();
        q.pop();
        
        long long x = p[i];
        if(c == 0){
            // x пошёл в A => (a - x) тоже должен пойти в A
            long long y = a - x;
            int j = indexOf[y]; // по условию есть
            // Заставляем j быть тоже в A
            if(!fixColor(j, 0)){
                cout << "NO\n";
                return 0;
            }
        } else {
            // x пошёл в B => (b - x) тоже должен пойти в B
            long long y = b - x;
            int j = indexOf[y];
            if(!fixColor(j, 1)){
                cout << "NO\n";
                return 0;
            }
        }
    }
    
    // Если ещё остались элементы possible[i] == 3 (оба варианта),
    // мы можем свободно выбрать им, например, A (color=0) – 
    // и если это где-то далее вызовет фиксации, тоже прогоняем.
    for(int i=0; i<n; i++){
        if(possible[i] == 3){
            // Попробуем зафиксировать в A (c=0),
            // это точно не даст противоречий, 
            // ведь если "можно и то, и другое", 
            // то выбирать A – законно.
            // Но важно запустить fixColor, чтобы цепочки дозаполнились.
            if(!fixColor(i, 0)){
                // если вдруг конфликт (что маловероятно,
                // т.к. мы проверили "possible[i] == 3"),
                // тогда была бы опция fixColor(i,1).
                // Но с учётом логики — не должно здесь быть конфликта.
                cout << "NO\n";
                return 0;
            }
            // Запустим цикл обработки, пока все цепочки не улягутся
            while(!q.empty()){
                auto [ii, cc] = q.front();
                q.pop();
                long long xx = p[ii];
                if(cc == 0){
                    long long yy = a - xx;
                    int jj = indexOf[yy];
                    if(!fixColor(jj,0)){
                        cout << "NO\n";
                        return 0;
                    }
                } else {
                    long long yy = b - xx;
                    int jj = indexOf[yy];
                    if(!fixColor(jj,1)){
                        cout << "NO\n";
                        return 0;
                    }
                }
            }
        }
    }
    
    // Если дошли сюда — решение есть. 
    cout << "YES\n";
    // Выведем раскладку в том же порядке, что и входные p[i].
    // У нас: possible[i] будет 1 (значит "A") или 2 (значит "B").
    // Нужно вывести 0 для A, 1 для B:
    for(int i=0; i<n; i++){
        if(possible[i] == 1){
            cout << 0 << " "; // A
        } else {
            cout << 1 << " "; // B
        }
    }
    cout << "\n";
    
    return 0;
}

Пояснения к коду:

• Мы сначала сохраняем в possible[i] все допустимые варианты для $p_i$.

  • Если $(a - p_i)$ нет в словаре, убираем возможность кластm $p_i$ в $A$.
  • Если $(b - p_i)$ нет в словаре, убираем возможность кластm $p_i$ в $B$.
  • Если после этого ничего не остаётся, сразу «NO».

• Далее все элементы, у которых осталась ровно одна возможность (т.е. possible[i] == 1 или possible[i] == 2), мы складываем в очередь q как жёстко зафиксированные.

• В цикле while(!q.empty()) берём «зафиксированный» элемент $x$. Если он зафиксирован в $A$, то и $(a-x)$ обязан быть в $A$. Если у $(a-x)$ изначально была возможность (B) — её теперь нужно исключить. Если окажется, что у $(a-x)$ вообще больше нет вариантов, значит противоречие. Если у $(a-x)$ после «обрезания» остался ровно один вариант, мы тоже добавляем $(a-x)$ в очередь и так рекурсивно/каскадно. Аналогично с $B$.

• После полного «пропихивания» очереди может остаться часть элементов, у которых possible[i] == 3, т.е. они по-прежнему «могут и в $A$, и в $B$». Для них мы делаем выбор (скажем, попробуем «в $A$»), и если это порождает новые фиксации, опять пойдёт очередь. Если возникнет противоречие, теоретически можно было бы попробовать «в $B$», но в данной задаче (как показывает анализ) конфликт в таком месте не возникает, поскольку possible[i] == 3 означает, что $(a-p_i)$ и $(b-p_i)$ оба лежат в наборе, и цепочки таких «двусторонних» вершин не приведут к тупику — мы всегда можем выбрать любую сторону, и всё согласуется (если не возникало обязательств от соседей).

• Если в итоге конфликтов не обнаружилось, мы печатаем «YES» и конечную раскладку.

Таким образом за один-полтора прохода (плюс очередь) всё раскладывается за $O(n)$. Если где-то по пути ловим противоречие — «NO».

Это ровно соответствует примеру, где во втором тесте ( $p=\{1,2,4\}, a=3, b=4$ ) моментально выяснится, что для элемента $4$ изначально possible[..]==0 (оба бита выкидываются), и мы печатаем «NO». А в первом тесте ( $p=\{2,3,4,5\}, a=5, b=9$ ) получаем ровное распределение.

Итог:

• Идея основная — учесть, что «$x$ в $A$» возможно лишь при наличии $a-x$ в нашем наборе. Аналогично для $B$.
• Затем «склеивать»/«протаскивать» эти требования через очередь или через «систему взаимных ограничений».
• Если в каком-то месте вариантов не остаётся, ответ «NO». Иначе «YES» + восстановленная раскладка.