kz-maxx

RU KZ
← Главная

Бор

Trie для строк: вставка, поиск, подсчет префиксов и реализация через массивы

Что такое бор?

Бор, или trie, — это дерево для хранения строк. В нем каждая вершина соответствует какому-то префиксу.

root — пустая строка
ребро — один символ
путь от root — строка или префикс строки

Бор удобен, когда нужно быстро работать со множеством строк:

  • добавлять строки
  • проверять, есть ли строка
  • проверять, есть ли префикс
  • считать количество строк с данным префиксом
  • искать слова в словаре

Пример

Пусть в бор добавлены строки: 

ab
abc
ac

Тогда общий префикс a будет храниться один раз.

root
└── a
    ├── b
    │   └── c
    └── c

Строка ab заканчивается в вершине b, строка abc заканчивается в вершине c, строка ac заканчивается в другой вершине c.

Главная идея

Когда вставляем строку, идем по символам слева направо. Если нужного перехода нет, создаем новую вершину.

текущая вершина v
текущий символ c
если перехода нет → создаем вершину
переходим в следующую вершину
в конце отмечаем конец слова

Реализация через массивы

Для маленьких английских букв a...z удобно использовать массив: 

int tr[N][26];

Здесь:

tr[v][c] — куда идем из вершины v по символу c
если tr[v][c] = 0, перехода нет
cnt[v] — сколько слов заканчивается в вершине v
pass[v] — сколько слов проходит через вершину v
idx — количество созданных вершин
Обычно root имеет номер 0. Новые вершины получают номера 1, 2, 3 и так далее.

Базовые массивы

const int N = 1000005;

int tr[N][26];
int cnt[N];
int pass[N];
int idx;

Если сумма длин всех строк не больше S, то вершин в боре будет не больше S + 1.

Поэтому размер N нужно брать примерно как сумма длин всех добавляемых строк + 5.

Как символ перевести в число?

Для маленьких английских букв: 

int c = s[i] - 'a';

Тогда:

'a' → 0
'b' → 1
'c' → 2
...
'z' → 25

Вставка строки

Функция insert добавляет строку в бор. 

void insert(string s) {
    int v = 0;

    pass[v]++;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            tr[v][c] = ++idx;
        }

        v = tr[v][c];
        pass[v]++;
    }

    cnt[v]++;
}

В конце:

cnt[v]++ означает, что в этой вершине заканчивается слово
pass[v]++ означает, что через эту вершину проходит слово

Поиск полной строки

Чтобы проверить, есть ли строка в боре, идем по символам строки.

Если где-то нет перехода, строки нет. Если дошли до конца, проверяем cnt[v]. 

bool search(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return false;
        }

        v = tr[v][c];
    }

    return cnt[v] > 0;
}

Проверка префикса

Чтобы проверить, есть ли хотя бы одно слово с таким префиксом, просто идем по символам префикса. 

bool startsWith(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return false;
        }

        v = tr[v][c];
    }

    return true;
}

Здесь не нужно проверять cnt[v], потому что нам важен не конец слова, а существование пути.

Сколько слов имеют данный префикс?

Для этого нужен массив pass. Он показывает, сколько слов проходит через вершину. 

int countPrefix(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return 0;
        }

        v = tr[v][c];
    }

    return pass[v];
}

Сколько раз добавили конкретное слово?

Для этого используем cnt[v]. Он хранит количество слов, которые заканчиваются в вершине v. 

int countWord(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return 0;
        }

        v = tr[v][c];
    }

    return cnt[v];
}

Полный код бора на массивах

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

const int N = 1000005;

int tr[N][26];
int cnt[N];
int pass[N];
int idx;

void insert(string s) {
    int v = 0;

    pass[v]++;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            tr[v][c] = ++idx;
        }

        v = tr[v][c];
        pass[v]++;
    }

    cnt[v]++;
}

bool search(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return false;
        }

        v = tr[v][c];
    }

    return cnt[v] > 0;
}

bool startsWith(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return false;
        }

        v = tr[v][c];
    }

    return true;
}

int countWord(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return 0;
        }

        v = tr[v][c];
    }

    return cnt[v];
}

int countPrefix(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return 0;
        }

        v = tr[v][c];
    }

    return pass[v];
}

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

    insert("ab");
    insert("abc");
    insert("ac");
    insert("ab");

    cout << search("ab") << '\n';
    cout << search("a") << '\n';
    cout << startsWith("a") << '\n';
    cout << countWord("ab") << '\n';
    cout << countPrefix("a") << '\n';

    return 0;
}

Output примера

1
0
1
2
4

Почему:

  • search("ab") = 1, потому что слово есть
  • search("a") = 0, потому что префикс есть, но слова a нет
  • startsWith("a") = 1, потому что есть слова на a
  • countWord("ab") = 2, потому что ab добавили два раза
  • countPrefix("a") = 4, потому что все 4 добавленных слова начинаются с a

Удаление слова

Если слово было добавлено, его можно удалить. Нужно пройти по пути и уменьшить pass, а в конце уменьшить cnt. 

bool eraseWord(string s) {
    if (countWord(s) == 0) {
        return false;
    }

    int v = 0;

    pass[v]--;

    for (char ch : s) {
        int c = ch - 'a';

        v = tr[v][c];
        pass[v]--;
    }

    cnt[v]--;

    return true;
}
В этой простой версии мы не удаляем вершины физически, просто уменьшаем счетчики.

Команды insert / find / prefix

Часто в задачах дают запросы:

add s — добавить строку
find s — проверить полную строку
pref s — сколько строк имеют префикс s

Полная программа с запросами

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

const int N = 1000005;

int tr[N][26];
int cnt[N];
int pass[N];
int idx;

void insert(string s) {
    int v = 0;

    pass[v]++;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            tr[v][c] = ++idx;
        }

        v = tr[v][c];
        pass[v]++;
    }

    cnt[v]++;
}

int countWord(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return 0;
        }

        v = tr[v][c];
    }

    return cnt[v];
}

int countPrefix(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return 0;
        }

        v = tr[v][c];
    }

    return pass[v];
}

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

    int q;
    cin >> q;

    while (q--) {
        string type, s;
        cin >> type >> s;

        if (type == "add") {
            insert(s);
        } else if (type == "find") {
            if (countWord(s) > 0) {
                cout << "YES\n";
            } else {
                cout << "NO\n";
            }
        } else if (type == "pref") {
            cout << countPrefix(s) << '\n';
        }
    }

    return 0;
}

Input / Output

Input

7
add ab
add abc
add ac
find ab
find a
pref a
pref ab

Output

YES
NO
3
2

Бор для цифр

Если строки состоят из цифр 0...9, нужно поменять размер алфавита на 10. 

int tr[N][10];

int c = ch - '0';

Это часто встречается в задачах про номера телефонов.

Бор для больших букв

Если есть только большие буквы A...Z: 

int tr[N][26];

int c = ch - 'A';

Если алфавит большой

Если символов много, массив tr[N][alphabet] может занимать слишком много памяти. Тогда можно использовать map или unordered_map. 

vector<unordered_map<char, int>> tr;
Но для маленьких английских букв массив быстрее и проще.

Проверка на конфликт телефонных номеров

Типичная задача: есть список телефонных номеров. Нужно проверить, является ли один номер префиксом другого.

Например: 

911
9112

Здесь 911 — префикс 9112, значит есть конфликт.

Идея проверки конфликта

Когда вставляем строку:

  • если во время движения встретили вершину, где уже заканчивается слово, значит старое слово является префиксом нового
  • если после вставки у текущей вершины есть дети, значит новое слово является префиксом старого

Код проверки конфликта

int tr[N][10];
int cnt[N];
int idx;

bool hasChild(int v) {
    for (int c = 0; c < 10; c++) {
        if (tr[v][c] != 0) {
            return true;
        }
    }

    return false;
}

bool insertPhone(string s) {
    int v = 0;

    for (char ch : s) {
        if (cnt[v] > 0) {
            return false;
        }

        int c = ch - '0';

        if (tr[v][c] == 0) {
            tr[v][c] = ++idx;
        }

        v = tr[v][c];
    }

    if (cnt[v] > 0) {
        return false;
    }

    if (hasChild(v)) {
        return false;
    }

    cnt[v]++;

    return true;
}
false означает, что найден конфликт префиксов.

Асимптотика

Операция Время
Вставить строку длины L O(L)
Найти строку длины L O(L)
Проверить префикс длины L O(L)
Посчитать слова с префиксом O(L)

Память:

O(количество вершин * размер алфавита)

Для tr[N][26] память примерно: 

N * 26 * sizeof(int)

Когда использовать бор?

нужно хранить много строк
нужно быстро искать строки
нужно быстро искать префиксы
нужно считать количество строк с префиксом
нужно проверять prefix conflict
нужен словарь строк

Типичные ошибки

  • Забыть, что root — это вершина 0
  • Использовать 0 как новую вершину, хотя 0 уже root
  • Не увеличить idx при создании вершины
  • Забыть отметить конец слова через cnt[v]++
  • Путать cnt и pass
  • Взять маленький размер массива N
  • Использовать tr[N][26] для цифр и делать ch - 'a'
  • Не обнулить массивы между тестами

Как очистить бор между тестами?

Если тестов много, можно очистить только использованные вершины. 

void clearTrie() {
    for (int i = 0; i <= idx; i++) {
        for (int c = 0; c < 26; c++) {
            tr[i][c] = 0;
        }

        cnt[i] = 0;
        pass[i] = 0;
    }

    idx = 0;
}
Это быстрее, чем каждый раз очищать весь огромный массив N.

Короткий шаблон

const int N = 1000005;

int tr[N][26];
int cnt[N];
int pass[N];
int idx;

void insert(string s) {
    int v = 0;

    pass[v]++;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            tr[v][c] = ++idx;
        }

        v = tr[v][c];
        pass[v]++;
    }

    cnt[v]++;
}

int countWord(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return 0;
        }

        v = tr[v][c];
    }

    return cnt[v];
}

int countPrefix(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            return 0;
        }

        v = tr[v][c];
    }

    return pass[v];
}

Главное запомнить

бор = дерево префиксов
root = 0
tr[v][c] = переход из вершины v по символу c
cnt[v] = сколько слов заканчивается в v
pass[v] = сколько слов проходит через v
insert/search/prefix работают за O(длина строки)