Разбор задачи «Межзвёздная экспедиция»

Идея задачи

У нас есть n пилотов и m экспедиций.

Для каждого пилота известна его устойчивость a_j, а для каждой экспедиции — сложность b_i.

Если на экспедицию назначено k пилотов, то каждый из них должен выдерживать не меньше b_i / k. Удобнее использовать эквивалентную формулировку:

если самый слабый пилот в группе имеет устойчивость x, то группа подходит для экспедиции тогда и только тогда, когда

x * k >= b_i.

Нужно либо построить распределение пилотов по экспедициям, либо понять, что это невозможно.

Главная сложность задачи в том, что пилотов много, до 2 * 10^5, а экспедиций мало, до 20.

Это почти всегда означает, что стоит искать решение через DP по подмножествам экспедиций.

Первое важное наблюдение

Отсортируем пилотов по убыванию устойчивости.

Пусть мы каким-то образом смогли закрыть несколько экспедиций, используя x пилотов. Тогда можно считать, что использованы именно x самых сильных пилотов.

Почему это верно?

Если в каком-то решении используется более слабый пилот, а более сильный при этом свободен, то можно просто заменить слабого на сильного. Условия ни для одной экспедиции не ухудшатся.

Значит, после сортировки достаточно рассматривать только такую картину:

• мы берем некоторый префикс самых сильных пилотов,
• разбиваем его на несколько групп,
• и эти группы назначаем на экспедиции.

Это очень сильное упрощение.

Второе важное наблюдение

Предположим, что мы уже использовали первых x самых сильных пилотов для части экспедиций.

Теперь хотим закрыть ещё одну экспедицию. Тогда выгоднее всего брать для неё следующих подряд пилотов:

x + 1, x + 2, ..., x + k.

Почему именно так?

Потому что среди всех ещё не использованных пилотов это самые сильные. Если какая-то группа размера k вообще может выполнить экспедицию, то группа из k самых сильных свободных пилотов тоже сможет.

Значит, каждая новая экспедиция в нашем DP будет получать подряд идущий блок пилотов.

Как проверить, подходит ли блок пилотов

Пусть мы рассматриваем экспедицию сложности b.

Мы хотим дать ей блок пилотов

x + 1, x + 2, ..., x + k.

Так как пилоты отсортированы по убыванию, самый слабый в этом блоке — это пилот на позиции x + k.

Тогда условие такое:

a[x + k] * k >= b

Нужно найти минимальное k, для которого это выполняется.

Если мы умеем быстро находить такой k, то дальше задача хорошо ложится на DP.

DP по маске экспедиций

Будем хранить:

• dp[mask] — минимальное число первых пилотов, которых достаточно, чтобы закрыть все экспедиции из множества mask.

Тогда:

• dp[0] = 0, потому что для пустого набора экспедиций никого не нужно;
• если dp[mask] = x, значит, экспедиции из mask можно закрыть, используя первых x пилотов;
• хотим добавить новую экспедицию i, которая пока не входит в mask;
• нужно понять, какое минимальное число следующих пилотов надо взять для этой экспедиции.

Если минимальный подходящий блок заканчивается на позиции r, то:

• эта экспедиция получает пилотов с позиций от x + 1 до r,
• новое состояние: dp[mask | (1 << i)] = min(dp[mask | (1 << i)], r).

Так как m <= 20, количество состояний равно 2^m, и это уже подходит.

Как быстро находить следующий подходящий блок

Наивно для каждого состояния и каждой экспедиции перебирать длину блока слишком медленно.

Нужно сделать предподсчёт.

Пусть мы хотим для экспедиции i узнать:

• если уже использовано x пилотов,
• где закончится минимальный подходящий следующий блок?

Обозначим это как nxt[i][x].

То есть:

• nxt[i][x] = r, если минимальный подходящий блок — это пилоты x + 1 ... r;
• если такого блока нет, будем считать значение бесконечным.

Как посчитать nxt

Зафиксируем экспедицию i и позицию r — кандидата на последнего пилота блока.

Тогда длина блока равна r - x, а самый слабый пилот — a[r].

Условие:

a[r] * (r - x) >= b[i]

Отсюда:

r - x >= ceil(b[i] / a[r])

то есть

x <= r - ceil(b[i] / a[r])

Это означает, что для фиксированного r подходят все значения x от 0 до некоторой границы.

Если пройти r слева направо, то можно заполнять ещё не заполненные позиции x, для которых текущий r уже подходит.

Чтобы делать это быстро, удобно использовать DSU «следующий непомеченный индекс».

Именно так и строится массив nxt.

Восстановление ответа

Кроме массива dp, будем хранить:

• parent_mask[new_mask] — из какого состояния пришли,
• parent_proj[new_mask] — какую экспедицию добавили.

Тогда после нахождения полного состояния можно пройти назад по родителям.

Если переход был из состояния pm в состояние mask, то экспедиция parent_proj[mask] получила пилотов с позиций:

dp[pm] + 1 ... dp[mask]

Так как мы запоминаем исходные номера пилотов, легко вывести ответ в требуемом формате.

Почему решение корректно

Утверждение 1

Если некоторый набор экспедиций можно закрыть x пилотами, то его можно закрыть первыми x самыми сильными пилотами.

Доказательство.

Возьмём любое корректное распределение. Если существует неиспользованный пилот сильнее какого-то использованного, то заменим использованного на него. Условия для всех экспедиций сохраняются или улучшаются. Повторяя этот процесс, получим корректное распределение, использующее именно x самых сильных пилотов.

Утверждение 2

Если уже использовано x пилотов, то для новой экспедиции достаточно рассматривать блок из следующих подряд пилотов.

Доказательство.

Среди всех свободных пилотов пилоты с позициями x + 1, x + 2, ... являются самыми сильными. Если существует подходящая группа размера k, то группа из k самых сильных свободных пилотов тоже будет подходящей. Значит, можно ограничиться подряд идущим блоком.

Утверждение 3

Переходы DP корректны.

Доказательство.

По первому утверждению любое допустимое частичное решение можно представить как использование некоторого префикса пилотов. По второму утверждению при добавлении новой экспедиции достаточно взять минимальный подходящий следующий блок. Именно это и делает переход DP.

Итог

Если dp[(1 << m) - 1] конечно, то мы построили корректное распределение пилотов по всем экспедициям. Если нет, то подходящего распределения не существует.

Сложность

Пусть m <= 20, n <= 2 * 10^5.

• Предподсчёт nxt: O(m * n * alpha(n))
• DP по маске: O(m * 2^m)
• Память: O(m * n + 2^m)

Этого достаточно для заданных ограничений.

Реализация на C++

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

struct DSU {
    vector<int> p;

    DSU() {}
    DSU(int n) { init(n); }

    void init(int n) {
        p.resize(n + 1);
        iota(p.begin(), p.end(), 0);
    }

    int find(int v) {
        if (p[v] == v) return v;
        return p[v] = find(p[v]);
    }

    void erase_pos(int v) {
        p[find(v)] = find(v + 1);
    }
};

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

    int n, m;
    cin >> n >> m;

    vector<long long> a(n + 1);
    vector<pair<long long, int>> pilots;
    pilots.reserve(n);

    for (int i = 1; i <= n; i++) {
        cin >> a[i];
        pilots.push_back({a[i], i});
    }

    vector<long long> b(m);
    for (int i = 0; i < m; i++) cin >> b[i];

    sort(pilots.begin(), pilots.end(), [&](auto &x, auto &y) {
        return x.first > y.first;
    });

    vector<long long> sortedA(n + 1);
    vector<int> origId(n + 1);

    for (int i = 1; i <= n; i++) {
        sortedA[i] = pilots[i - 1].first;
        origId[i] = pilots[i - 1].second;
    }

    const int INF = (int)1e9;
    vector<vector<int>> nxt(m, vector<int>(n + 1, INF));

    for (int proj = 0; proj < m; proj++) {
        DSU dsu(n);

        for (int r = 1; r <= n; r++) {
            long long need = (b[proj] + sortedA[r] - 1) / sortedA[r];
            long long bound = 1LL * r - need;

            if (bound < 0) continue;
            if (bound > r - 1) bound = r - 1;

            int x = dsu.find(0);
            while (x <= bound) {
                nxt[proj][x] = r;
                dsu.erase_pos(x);
                x = dsu.find(x);
            }
        }
    }

    int FULL = 1 << m;
    vector<int> dp(FULL, INF);
    vector<int> parentMask(FULL, -1);
    vector<int> parentProj(FULL, -1);

    dp[0] = 0;

    for (int mask = 0; mask < FULL; mask++) {
        if (dp[mask] == INF) continue;
        int used = dp[mask];
        if (used > n) continue;

        for (int proj = 0; proj < m; proj++) {
            if (mask & (1 << proj)) continue;

            int r = nxt[proj][used];
            if (r == INF) continue;

            int newMask = mask | (1 << proj);
            if (r < dp[newMask]) {
                dp[newMask] = r;
                parentMask[newMask] = mask;
                parentProj[newMask] = proj;
            }
        }
    }

    if (dp[FULL - 1] == INF || dp[FULL - 1] > n) {
        cout << "NO\n";
        return 0;
    }

    vector<vector<int>> ans(m);
    int mask = FULL - 1;

    while (mask != 0) {
        int pm = parentMask[mask];
        int proj = parentProj[mask];

        int l = dp[pm] + 1;
        int r = dp[mask];

        for (int pos = l; pos <= r; pos++) {
            ans[proj].push_back(origId[pos]);
        }

        mask = pm;
    }

    cout << "YES\n";
    for (int i = 0; i < m; i++) {
        cout << ans[i].size();
        for (int id : ans[i]) cout << ' ' << id;
        cout << '\n';
    }

    return 0;
}

Реализация на Python

import sys

INF = 10**18


class DSU:
    def __init__(self, n):
        self.p = list(range(n + 1))

    def find(self, v):
        while self.p[v] != v:
            self.p[v] = self.p[self.p[v]]
            v = self.p[v]
        return v

    def erase_pos(self, v):
        self.p[self.find(v)] = self.find(v + 1)


def solve():
    input = sys.stdin.readline

    n, m = map(int, input().split())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))

    pilots = [(a[i], i + 1) for i in range(n)]
    pilots.sort(reverse=True)

    sorted_a = [0] * (n + 1)
    orig_id = [0] * (n + 1)

    for i in range(1, n + 1):
        sorted_a[i] = pilots[i - 1][0]
        orig_id[i] = pilots[i - 1][1]

    nxt = [[INF] * (n + 1) for _ in range(m)]

    for proj in range(m):
        dsu = DSU(n)
        need_b = b[proj]

        for r in range(1, n + 1):
            need = (need_b + sorted_a[r] - 1) // sorted_a[r]
            bound = r - need

            if bound < 0:
                continue
            if bound > r - 1:
                bound = r - 1

            x = dsu.find(0)
            while x <= bound:
                nxt[proj][x] = r
                dsu.erase_pos(x)
                x = dsu.find(x)

    full = 1 << m
    dp = [INF] * full
    parent_mask = [-1] * full
    parent_proj = [-1] * full

    dp[0] = 0

    for mask in range(full):
        if dp[mask] == INF:
            continue

        used = dp[mask]
        if used > n:
            continue

        for proj in range(m):
            if (mask >> proj) & 1:
                continue

            r = nxt[proj][used]
            if r == INF:
                continue

            new_mask = mask | (1 << proj)
            if r < dp[new_mask]:
                dp[new_mask] = r
                parent_mask[new_mask] = mask
                parent_proj[new_mask] = proj

    if dp[full - 1] == INF or dp[full - 1] > n:
        print("NO")
        return

    ans = [[] for _ in range(m)]
    mask = full - 1

    while mask != 0:
        pm = parent_mask[mask]
        proj = parent_proj[mask]

        l = dp[pm] + 1
        r = dp[mask]

        for pos in range(l, r + 1):
            ans[proj].append(orig_id[pos])

        mask = pm

    print("YES")
    for group in ans:
        print(len(group), *group)


if __name__ == "__main__":
    solve()

Что важно запомнить из этой задачи

1. Когда объектов очень много, а групп или подзадач мало, часто помогает DP по маске.
2. Перед DP полезно искать структурное упрощение. Здесь им стала идея, что можно работать только с префиксом самых сильных пилотов.
3. После сортировки задача свелась к разбиению префикса на подряд идущие блоки.
4. Предподсчёт nxt позволил быстро делать переходы между состояниями.

Это очень хороший пример задачи, где основная трудность не в самой реализации DP, а в том, чтобы сначала правильно увидеть структуру ответа.