Запросы к бору

Бор — не только хранилище. По нему отвечают на те же запросы, что и по упорядоченному множеству, и делают это за длину строки, а не за логарифм от их числа.

Всюду ниже используется счётчик sub — число строк в поддереве вершины. Без него ни один из запросов не работает корректно после удалений.

Удаление

void erase(const string& s) {
    if (!contains(s)) return;                 // иначе счётчики уйдут в минус
    int v = 0;
    t[v].sub--;
    for (char ch : s) { v = t[v].next[ch - 'a']; t[v].sub--; }
    t[v].cnt--;
}

Вершины физически не удаляются — только уменьшаются счётчики. Освобождать память смысла нет: она всё равно уже выделена, а перестройка дерева стоила бы дорого.

Проверка contains перед удалением обязательна. Удаление отсутствующей строки испортит sub на всём пути, и дальше сломается всё остальное.

Минимальная строка, не меньшая данной

Аналог lower_bound. Алгоритм спускается по бору и в каждой вершине делает выбор:

1. Попробовать пойти по тому же символу, что в запросе, и решить ту же задачу ниже.
2. Если там ответа не нашлось — пойти по минимальному большему символу и дальше спускаться по самым маленьким рёбрам до первой терминальной вершины.
3. Если и такого нет — ответа в этом поддереве нет.

bool descendMin(int v, string& out) {         // самая маленькая строка поддерева
    if (t[v].sub == 0) return false;
    if (t[v].cnt > 0) return true;
    for (int c = 0; c < ALPHA; c++) {
        int u = t[v].next[c];
        if (u != -1 && t[u].sub > 0) { out += 'a' + c; return descendMin(u, out); }
    }
    return false;
}

bool go(int v, const string& s, size_t i, string& out) {
    if (t[v].sub == 0) return false;          // пустое поддерево — сразу мимо
    if (i == s.size()) return descendMin(v, out);
    int c = s[i] - 'a';
    int u = t[v].next[c];
    if (u != -1 && t[u].sub > 0) {            // ветка 1: точное совпадение символа
        size_t keep = out.size();
        out += s[i];
        if (go(u, s, i + 1, out)) return true;
        out.resize(keep);                     // не получилось — откатываем
    }
    for (int d = c + 1; d < ALPHA; d++) {     // ветка 2: первый больший символ
        int w = t[v].next[d];
        if (w != -1 && t[w].sub > 0) { out += 'a' + d; return descendMin(w, out); }
    }
    return false;
}

Проверено: 40 000 сценариев по 25 операций вперемешку (вставка, удаление, запрос) — ответы совпали с multiset<string>.

Строка if (t[v].sub == 0) return false

Вот ради чего заводился счётчик. Без него после удалений в боре остаются вершины и рёбра, под которыми нет ни одной строки — и ветка 2 радостно спустится в такое поддерево, не найдя там терминальной вершины.

Конкретный сценарий: добавили aaa, aab, aba, потом всё удалили, ищем строку не меньше aac. По символу a идём, по a идём, по c перехода нет — ветка 1 провалилась. Возвращаемся, ищем больший символ, находим b, спускаемся — и упираемся в конец пути без терминальной вершины. Алгоритм либо вернёт мусор, либо упадёт.

Со счётчиком такое поддерево отсекается сразу.

Сколько это стоит

Утверждение: $O(|s| + |\text{ответ}|)$, плюс множитель $|\Sigma|$ на перебор символов в вершине.

Ключ к оценке — мы никогда не возвращаемся в ветку 1 дважды. Пока символы совпадают, идём вниз по одному пути. Как только совпадение прервалось, уходим в ветку 2, а она уже не откатывается: если в поддереве есть хоть одна строка (а sub > 0 это гарантирует), то descendMin её найдёт. Дальше только спуск.

Без проверки sub ветка 2 могла бы провалиться, и рекурсия начала бы честно перебирать варианты — тогда оценка ломается вместе с корректностью.

$k$-я строка по порядку

Тот же счётчик даёт порядковую статистику. Спускаемся от корня, перебирая символы по возрастанию: если в поддереве очередного ребра строк меньше, чем $k$, вычитаем их число и идём к следующему символу.

string kth(int k) {                            // k от 1 до t[0].sub
    string res;
    int v = 0;
    while (true) {
        if (t[v].cnt >= k) return res;         // строки, кончающиеся здесь, идут первыми
        k -= t[v].cnt;
        for (int c = 0; c < ALPHA; c++) {
            int u = t[v].next[c];
            if (u == -1 || t[u].sub == 0) continue;
            if (k <= t[u].sub) { res += 'a' + c; v = u; break; }
            k -= t[u].sub;
        }
    }
}

Порядок проверок важен: строки, кончающиеся в самой вершине, лексикографически меньше всех, что лежат ниже, поэтому cnt списывается первым.

Тот же приём — спуск с вычитанием размеров поддеревьев — работает в любом дереве, где известны размеры: так же ищут $k$-й элемент в декартовом дереве и в дереве отрезков по префиксным суммам.

Сколько строк начинается на данный префикс

Самый дешёвый запрос: пройти по префиксу и вернуть sub конечной вершины. $O(|p|)$, никакого перебора.

Если вершины на пути нет — ответ ноль.