Aho Corasick

Trie và liên kết thất bại

Chèn mọi mẫu vào trie; nút $u$ ứng với xâu $str(u)$ là một tiền tố của mẫu nào đó. Liên kết thất bại fail[u] trỏ tới nút ứng với hậu tố thật dài nhất của $str(u)$ cũng có mặt trong trie. Tính bằng BFS theo độ sâu: con $v = $ nxt[u][c] có fail[v] = nxt[fail[u]][c], vì hậu tố dài nhất của $str(u) + c$ có trong trie là một hậu tố của $str(u)$ (thứ có trong trie) nối thêm $c$.

Cùng lúc, lấp các ô chuyển còn trống: nếu nxt[u][c] không có thì đặt bằng nxt[fail[u]][c]. Khi đó ôtômat có bảng chuyển đầy đủ, quét văn bản chỉ là u = nxt[u][c] mỗi ký tự, $O(1)$, không có vòng lặp lùi như KMP. Tổng $O(\Sigma \cdot \text{số nút})$ cho việc dựng.

Đếm xuất hiện. Ở trạng thái $u$ sau khi đọc tới vị trí $i$, các mẫu kết thúc tại $i$ là các mẫu ứng với $u$ và mọi tổ tiên của $u$ trên cây liên kết thất bại. Vậy đánh dấu cnt[u]++ khi đi qua, rồi cộng dồn từ con lên cha theo thứ tự BFS ngược: cnt[fail[u]] += cnt[u]. Số lần xuất hiện của mẫu là cnt tại nút cuối của nó.

Ví dụSổ Đếm Mẫu

Đề. Văn bản dài tới $10^5$, tới $5 \cdot 10^5$ mẫu với tổng độ dài tới $5 \cdot 10^5$. Đếm số lần xuất hiện (được chồng lấn) của từng mẫu.

Lời giải. Chèn mọi mẫu vào trie, dựng liên kết thất bại và bảng chuyển đầy đủ bằng BFS. Quét văn bản: mỗi ký tự một bước chuyển, và tăng cnt của trạng thái vừa tới. Mẫu kết thúc tại vị trí $i$ là các nút trên chuỗi liên kết thất bại từ trạng thái ở $i$ lên gốc, nên số lần xuất hiện của mẫu kết thúc ở nút $x$ là tổng cnt trong cây con của $x$ trên cây liên kết thất bại.

Bước 1: chèn mẫu. Thiếu cạnh thì tạo nút mới; trả về nút cuối của mẫu. Hai mẫu giống nhau có cùng nút cuối và tự động có cùng đáp án.

int insert(string &p) {
    int u = 0;
    for (char ch : p) {
        int c = ch - 'a';
        if (!nxt[u][c]) nxt[u][c] = ++tot;
        u = nxt[u][c];
    }
    return u;
}

Bước 2: BFS dựng liên kết. BFS theo độ sâu nên fail[u] (nông hơn $u$) đã xong trước khi xét con của $u$. Có cạnh $u \xrightarrow{c} v$ thì fail[v] = nxt[fail[u]][c]; không có thì lấp nxt[u][c] = nxt[fail[u]][c]. Con của gốc có fail = 0 sẵn. Sau bước này nxt là bảng chuyển đầy đủ, quét văn bản chỉ là u = nxt[u][c], không có vòng lùi như KMP.

for (int c = 0; c < 26; c++) {
    int v = nxt[u][c];
    if (v) fail[v] = nxt[fail[u]][c], q.push(v);
    else nxt[u][c] = nxt[fail[u]][c];
}

Bước 3: quét và cộng dồn. order là thứ tự BFS; mỗi nút đứng sau fail của nó. Duyệt ngược thì mọi nút được cộng vào fail sau khi chính nó đã nhận đủ.

for (char ch : s) u = nxt[u][ch - 'a'], cnt[u]++;
for (int i = (int)order.size() - 1; i >= 0; i--) cnt[fail[order[i]]] += cnt[order[i]];

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

const int N = 5e5 + 5;
int tot, nxt[N][26], fail[N], cnt[N], endNode[N];
vector<int> order;

int insert(string &p) {
    int u = 0;
    for (char ch : p) {
        int c = ch - 'a';
        if (!nxt[u][c]) nxt[u][c] = ++tot;
        u = nxt[u][c];
    }
    return u;
}

void build() {
    queue<int> q;
    for (int c = 0; c < 26; c++)
        if (nxt[0][c]) q.push(nxt[0][c]);
    while (!q.empty()) {
        int u = q.front();
        q.pop();
        order.push_back(u);
        for (int c = 0; c < 26; c++) {
            int v = nxt[u][c];
            if (v) fail[v] = nxt[fail[u]][c], q.push(v);
            else nxt[u][c] = nxt[fail[u]][c];
        }
    }
}

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    string s;
    int k;
    cin >> s >> k;
    for (int i = 0; i < k; i++) {
        string p;
        cin >> p;
        endNode[i] = insert(p);
    }
    build();
    int u = 0;
    for (char ch : s) u = nxt[u][ch - 'a'], cnt[u]++;
    for (int i = (int)order.size() - 1; i >= 0; i--) cnt[fail[order[i]]] += cnt[order[i]];
    for (int i = 0; i < k; i++) cout << cnt[endNode[i]] << '\n';
}

ChạyInputaaaa
3
a
aa
aaa

Output4
3
2

$O(26 \cdot \text{số nút} + |S|)$.

Lan truyền trên cây liên kết thất bại

Mọi câu hỏi kiểu “thông tin của các vị trí mà mẫu kết thúc” đều đi theo một khuôn: ghi thông tin vào nút trạng thái tại từng vị trí của văn bản, rồi gộp từ con lên cha trên cây liên kết thất bại. Phép gộp có thể là cộng (số lần), lấy nhỏ nhất (vị trí kết thúc đầu tiên), lấy lớn nhất, hay OR (đã xuất hiện chưa). Mỗi nút chỉ gộp một lần nên tổng $O(\text{số nút})$, thay vì đi theo chuỗi liên kết thất bại ở mỗi vị trí văn bản, có thể tới $O(|S| \cdot \text{độ dài mẫu})$.

Ví dụVị Trí Đầu Tiên

Đề. Như bài trước, nhưng in vị trí bắt đầu nhỏ nhất (đánh số từ 1) của từng mẫu, hoặc $-1$.

Lời giải. Cùng khung với ví dụ trước (insert, build giống hệt), thay phép cộng bằng phép lấy nhỏ nhất. first[u] là vị trí kết thúc nhỏ nhất mà trạng thái bằng $u$. Gộp từ con lên cha trên cây liên kết thất bại cho vị trí kết thúc sớm nhất của mẫu ở mỗi nút. Vị trí bắt đầu là vị trí kết thúc trừ độ dài mẫu cộng 1.

Bước 1: quét. first khởi tạo INF; trạng thái ở vị trí $i$ (đánh số từ 1) nhận min(first[u], i + 1). Quét từ trái sang nên lần gán đầu tiên đã là nhỏ nhất.

for (int i = 0; i < (int)s.size(); i++) {
    u = nxt[u][s[i] - 'a'];
    first[u] = min(first[u], i + 1);
}

Bước 2: gộp lấy nhỏ nhất. Mẫu chưa từng xuất hiện giữ INF và in $-1$.

for (int i = (int)order.size() - 1; i >= 0; i--) {
    int v = order[i];
    first[fail[v]] = min(first[fail[v]], first[v]);
}

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

const int N = 5e5 + 5, INF = 1e9;
int tot, nxt[N][26], fail[N], first[N], endNode[N], len[N];
vector<int> order;

int insert(string &p) {
    int u = 0;
    for (char ch : p) {
        int c = ch - 'a';
        if (!nxt[u][c]) nxt[u][c] = ++tot;
        u = nxt[u][c];
    }
    return u;
}

void build() {
    queue<int> q;
    for (int c = 0; c < 26; c++)
        if (nxt[0][c]) q.push(nxt[0][c]);
    while (!q.empty()) {
        int u = q.front();
        q.pop();
        order.push_back(u);
        for (int c = 0; c < 26; c++) {
            int v = nxt[u][c];
            if (v) fail[v] = nxt[fail[u]][c], q.push(v);
            else nxt[u][c] = nxt[fail[u]][c];
        }
    }
}

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    string s;
    int k;
    cin >> s >> k;
    for (int i = 0; i < k; i++) {
        string p;
        cin >> p;
        endNode[i] = insert(p);
        len[i] = p.size();
    }
    build();
    for (int u = 0; u <= tot; u++) first[u] = INF;
    int u = 0;
    for (int i = 0; i < (int)s.size(); i++) {
        u = nxt[u][s[i] - 'a'];
        first[u] = min(first[u], i + 1);
    }
    for (int i = (int)order.size() - 1; i >= 0; i--) {
        int v = order[i];
        first[fail[v]] = min(first[fail[v]], first[v]);
    }
    for (int i = 0; i < k; i++) {
        int e = first[endNode[i]];
        cout << (e == INF ? -1 : e - len[i] + 1) << '\n';
    }
}

ChạyInputcababa
4
aba
cab
bb
a

Output2
1
-1
2

$O(26 \cdot \text{số nút} + |S|)$.

Liên kết đầu ra: các mẫu kết thúc tại một vị trí

Khi cần duyệt các mẫu kết thúc tại vị trí $i$ (không chỉ gộp), đi theo chuỗi liên kết thất bại có thể qua nhiều nút không phải cuối mẫu. Liên kết đầu ra out[u] trỏ thẳng tới tổ tiên gần nhất trên cây liên kết thất bại là cuối một mẫu: out[v] = term[fail[v]] ? fail[v] : out[fail[v]], tính trong BFS. Duyệt term[u] ? u : out[u], rồi out liên tiếp, chỉ đi qua đúng các mẫu kết thúc tại $i$.

Số mẫu phân biệt kết thúc tại một vị trí không vượt số độ dài khác nhau của mẫu, tức $O(\sqrt{L})$ với $L$ là tổng độ dài. Đây là chặn cho các DP ghép xâu.

Ví dụGhép Chuỗi Bằng Từ Điển

Đề. Xâu đích dài tới 5000, từ điển $k \le 10^5$ từ phân biệt, tổng độ dài tới $10^6$. Đếm số cách chia xâu đích thành dãy từ của từ điển, theo mô-đun $10^9 + 7$.

Lời giải. $dp_i$ là số cách ghép $i$ ký tự đầu, $dp_0 = 1$. Từ cuối cùng của một cách ghép kết thúc ở vị trí $i$, nên $dp_i = \sum dp_{i - l}$ trên mọi từ độ dài $l$ kết thúc tại $i$. Quét xâu đích bằng ôtômat Aho-Corasick; ở mỗi vị trí, liên kết đầu ra liệt kê đúng các từ kết thúc ở đó. Các từ đó có độ dài khác nhau, nên tối đa khoảng $\sqrt{2 \cdot 10^6} \approx 1400$ từ mỗi vị trí.

Bước 1: chèn có độ dài. dep[u] là độ dài xâu ứng với nút $u$, cũng là độ dài từ khi $u$ là nút cuối của một từ. Từ dài hơn xâu đích không bao giờ dùng được, bỏ ngay khi đọc để trie nhỏ đi.

Bước 2: liên kết đầu ra. out[v] là nút gần nhất trên chuỗi liên kết thất bại của $v$ (không kể $v$) là cuối một từ. Tính trong BFS vì fail[v] và out[fail[v]] đều đã có. Đi theo fail thay vì out cũng đúng, nhưng có thể qua rất nhiều nút không phải cuối từ ở mỗi vị trí.

fail[v] = nxt[fail[u]][c];
out[v] = term[fail[v]] ? fail[v] : out[fail[v]];

Bước 3: quy hoạch động. Bắt đầu ở $u$ nếu $u$ là cuối một từ, ngược lại ở out[u], rồi đi theo out tới 0. Mỗi nút gặp được là một từ kết thúc tại $i$.

for (int i = 1; i <= (int)s.size(); i++) {
    u = nxt[u][s[i - 1] - 'a'];
    for (int v = term[u] ? u : out[u]; v; v = out[v]) dp[i] = (dp[i] + dp[i - dep[v]]) % MOD;
}

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

const int N = 1e6 + 5;
const ll MOD = 1e9 + 7;
int tot, nxt[N][26], fail[N], out[N], dep[N];
bool term[N];
ll dp[5005];

void insert(string &p) {
    int u = 0;
    for (char ch : p) {
        int c = ch - 'a';
        if (!nxt[u][c]) nxt[u][c] = ++tot, dep[tot] = dep[u] + 1;
        u = nxt[u][c];
    }
    term[u] = 1;
}

void build() {
    queue<int> q;
    for (int c = 0; c < 26; c++)
        if (nxt[0][c]) q.push(nxt[0][c]);
    while (!q.empty()) {
        int u = q.front();
        q.pop();
        for (int c = 0; c < 26; c++) {
            int v = nxt[u][c];
            if (v) {
                fail[v] = nxt[fail[u]][c];
                out[v] = term[fail[v]] ? fail[v] : out[fail[v]];
                q.push(v);
            } else nxt[u][c] = nxt[fail[u]][c];
        }
    }
}

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    string s;
    int k;
    cin >> s >> k;
    for (int i = 0; i < k; i++) {
        string p;
        cin >> p;
        if (p.size() <= s.size()) insert(p);
    }
    build();
    dp[0] = 1;
    int u = 0;
    for (int i = 1; i <= (int)s.size(); i++) {
        u = nxt[u][s[i - 1] - 'a'];
        for (int v = term[u] ? u : out[u]; v; v = out[v]) dp[i] = (dp[i] + dp[i - dep[v]]) % MOD;
    }
    cout << dp[s.size()] << '\n';
}

ChạyInputba
3
b
a
ba

Output2

$O(26 \cdot \text{số nút} + |S| \sqrt{L})$ với $L$ là tổng độ dài từ điển.

Quy hoạch động trên ôtômat

Xâu không chứa mẫu cấm nào khi và chỉ khi khi đọc nó, ôtômat không bao giờ vào một nút “xấu”: nút là cuối một mẫu, hoặc có tổ tiên trên cây liên kết thất bại là cuối một mẫu (tính bad[v] |= bad[fail[v]] trong BFS). Đếm xâu dài $n$ tránh nút xấu là đếm đường đi dài $n$ trên đồ thị chuyển chỉ gồm nút tốt: DP $O(n \cdot \text{số nút} \cdot \Sigma)$, hoặc lũy thừa ma trận chuyển khi $n$ rất lớn.

Cùng khung cho: DP xác suất tránh mẫu, sửa ít ký tự nhất để tránh mẫu, điểm lớn nhất (thay phép cộng bằng cực đại trong nhân ma trận), đếm xâu chứa đủ một tập từ khóa (thêm mặt nạ bit vào trạng thái).

Ví dụChuỗi DNA Không Bệnh

Đề. Tới 10 mẫu cấm trên bảng chữ ACGT, mỗi mẫu dài tới 10. Đếm xâu DNA dài $n \le 2 \cdot 10^9$ không chứa mẫu cấm nào, theo mô-đun $10^5$.

Lời giải. Xâu không chứa mẫu cấm khi và chỉ khi đọc nó bằng ôtômat không bao giờ vào nút xấu: nút là cuối một mẫu, hoặc có tổ tiên trên cây liên kết thất bại là cuối một mẫu. Đếm xâu dài $n$ là đếm đường đi dài $n$ từ gốc trên đồ thị chuyển chỉ gồm nút tốt. Trie có tối đa 101 nút; ma trận $M_{uv}$ là số ký tự đưa nút tốt $u$ sang nút tốt $v$. Đáp án là tổng hàng 0 của $M^n$, tính bằng lũy thừa nhanh vì $n$ tới $2 \cdot 10^9$.

Bước 1: nút xấu lan theo liên kết. Chèn mẫu với code đổi ACGT thành $0..3$; nút cuối mẫu có bad = 1. Trong BFS thêm bad[v] = bad[v] || bad[fail[v]]; BFS đi theo độ sâu nên bad[fail[v]] đã đầy đủ khi xét $v$.

if (v) fail[v] = nxt[fail[u]][c], bad[v] = bad[v] || bad[fail[v]], q.push(v);
else nxt[u][c] = nxt[fail[u]][c];

Bước 2: ma trận chuyển. Từ mỗi nút tốt, mỗi ký tự đi tới một nút; nút đó tốt thì cộng 1 vào ô tương ứng. Hàng của nút xấu để toàn 0. res bắt đầu là ma trận đơn vị.

for (int u = 0; u < sz; u++) {
    res.a[u][u] = 1;
    if (bad[u]) continue;
    for (int c = 0; c < 4; c++)
        if (!bad[nxt[u][c]]) base.a[u][nxt[u][c]]++;
}

Bước 3: lũy thừa nhanh. Phần tử nhỏ hơn $10^5$, tích hai phần tử dưới $10^{10}$, cộng 101 tích vẫn vừa long long, nên mul chỉ lấy dư một lần cho mỗi ô.

while (n) {
    if (n & 1) res = mul(res, base);
    base = mul(base, base);
    n >>= 1;
}

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

const ll MOD = 100000;
int tot, nxt[105][4], fail[105];
bool bad[105];
int sz;

struct Mat {
    ll a[105][105];
};

Mat mul(Mat &x, Mat &y) {
    Mat z;
    for (int i = 0; i < sz; i++)
        for (int j = 0; j < sz; j++) {
            ll s = 0;
            for (int k = 0; k < sz; k++) s += x.a[i][k] * y.a[k][j];
            z.a[i][j] = s % MOD;
        }
    return z;
}

int code(char c) {
    return c == 'A' ? 0 : c == 'C' ? 1 : c == 'G' ? 2 : 3;
}

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    int m;
    ll n;
    cin >> m >> n;
    for (int i = 0; i < m; i++) {
        string p;
        cin >> p;
        int u = 0;
        for (char ch : p) {
            int c = code(ch);
            if (!nxt[u][c]) nxt[u][c] = ++tot;
            u = nxt[u][c];
        }
        bad[u] = 1;
    }
    queue<int> q;
    for (int c = 0; c < 4; c++)
        if (nxt[0][c]) q.push(nxt[0][c]);
    while (!q.empty()) {
        int u = q.front();
        q.pop();
        for (int c = 0; c < 4; c++) {
            int v = nxt[u][c];
            if (v) fail[v] = nxt[fail[u]][c], bad[v] = bad[v] || bad[fail[v]], q.push(v);
            else nxt[u][c] = nxt[fail[u]][c];
        }
    }
    sz = tot + 1;
    static Mat base, res;
    for (int i = 0; i < sz; i++)
        for (int j = 0; j < sz; j++) base.a[i][j] = res.a[i][j] = 0;
    for (int u = 0; u < sz; u++) {
        res.a[u][u] = 1;
        if (bad[u]) continue;
        for (int c = 0; c < 4; c++)
            if (!bad[nxt[u][c]]) base.a[u][nxt[u][c]]++;
    }
    while (n) {
        if (n & 1) res = mul(res, base);
        base = mul(base, base);
        n >>= 1;
    }
    ll ans = 0;
    for (int v = 0; v < sz; v++) ans += res.a[0][v];
    cout << ans % MOD << '\n';
}

ChạyInput2 2
A
CG

Output8

$O(\text{số nút}^3 \log n)$, khoảng $10^6 \cdot 31$ phép nhân.