Competitive Coding Algorithms

Dijkstra

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
int dijkstraM(int start, int end, vector<vector<pi>> adj, vector<int> &dist) {
    vector<bool> visited(N);

    priority_queue<pi, vector<pi>, greater<pi>> pq;

    pq.push({0, start});
    dist[start] = 0;

    while (pq.size()) {
        pi value = pq.top(); pq.pop();
        ll node = value.s;
        if (visited[node]) {
            continue;
        }
        visited[node] = true;

        for (auto to : adj[node]) {
            int newWeight = to.f + dist[node];

            if (newWeight < dist[to.s]) {
                // nodeInfo[to.s] = nodeInfo[node] + cows[to.s];
                prevV[to.s] = node;

                dist[to.s] = newWeight;
                pq.push({dist[to.s], to.s});
            }
        }

        //dbg(node, dist[node]);
    }

    return dist[end];
}

BIT

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
struct BIT {
    vector<int> tree;
    void init(int _n) { tree.resize(_n); }

    void upd(int i, int x) {
        while (i < tree.size()) {
            tree[i] = (tree[i] + x);
            //if (tree[i] < 0) tree[i] += M;
            if (tree[i] > M) tree[i] -= M;
            tree[i] = (tree[i] + M) % M;
            i += (i & (-i));
        }
    }

    long long query(int i) {
        long long ans = 0;
        while (i > 0) {
            ans += tree[i];
            //if (ans < 0) ans += M;
            if (ans > M) ans -= M;
            ans = (ans + M) % M;
            i -= (i & (-i));
        }
        return ans;
    }
};

BIT  allColors;
allColors.query(j);
allColors.upd(j, x);

Normal Segment Tree

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
vector<int> arr;
vector<int> tree;
arr.resize(N + 1);
tree.resize(4*N);

void build(int p, int l, int r) {
    if (l == r) {tree[p] = arr[l]; return;}

    int mid = (l + r)/2;
    build(2*p, l, mid);
    build(2*p + 1, mid + 1, r);

    tree[p] = min(tree[2*p], tree[2*p + 1]);
}

void update(int p, int l, int r, int x, int y) {
    // turn arr[x] into y
    if (l == r) { tree[p] = y; return; }

    int mid = (l + r)/2;
    if (x <= mid) {
        update(2*p, l, mid, x, y);
    }
    else {
        update(2*p + 1, mid + 1, r, x, y);
    }

    tree[p] = min(tree[2*p], tree[2*p + 1]);
}

int query(int p, int l, int r, int x, int y) {
    if (l == r) { return tree[p]; }

    int mid = (l + r)/2;
    int ans = INT_MAX;
    if (x <= mid) {
        ans = min(ans, query(2 * p, l, mid, x, y));
    }
    if (y >= mid + 1) {
        ans = min(ans, query(2 * p + 1, mid + 1, r, x, y));
    }

    return ans;
}

Lazy Segment Tree

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
template<class T, int MX> struct LazySeg {
    T treesum[MX], treemin[MX], lazy[MX];

    void build(int n, int tl, int tr) {
        if (tl == tr) {
            treesum[n] = arr[tl];
            treemin[n] = arr[tl];
            return;
        }
        int mid = (tl + tr)/2;
        build(n << 1, tl, mid);
        build((n << 1) + 1, mid + 1, tr);
        // change me
        treesum[n] = (treesum[n << 1] + treesum[(n << 1) + 1]);
        treemin[n] = min(treemin[2*n], treemin[2*n + 1]);
    }

    void pushdown(int n, int tl, int tr) {
        if (lazy[n] > 0) {
            // change me
            treesum[n] += (tr - tl + 1) * lazy[n];
            treemin[n] += lazy[n];

            if (tl != tr) {
                lazy[2*n] += lazy[n];
                lazy[2*n+1] += lazy[n];
            }
            lazy[n] = 0;
        }
    }

    void upd(int n, int tl, int tr, int st, int ed, T v) {
        pushdown(n, tl, tr);
        if (tl > tr || tl > ed || tr < st) return;
        if (tl >= st && tr <= ed) {
            // change me: only for min/max
            treesum[n] += (tr - tl + 1) * v;
            treemin[n] += v;

            // only for min
            if (tl != tr) {
                lazy[2*n] += v;
                lazy[2*n + 1] += v;
            }
            return;
        }
        else {
            int mid = (tl + tr)/2;
            upd(2*n, tl, mid, st, ed, v);
            upd(2*n+1, mid+1, tr, st, ed, v);

            // change me
            treesum[n] = (treesum[2*n] + treesum[2*n+1]);
            treemin[n] = min(treemin[2*n+1], treemin[2*n]);
        }
    }

    T querySum(int n, int tl, int tr, int st, int ed) {
        // change me
        if (tl > tr || ed < tl || st > tr) return 0;

        pushdown(n, tl, tr);

        // segment completely inside query
        if (st <= tl && tr <= ed) {
            return treesum[n];
        }
        else {
            int mid = (tl + tr)/2;
            T a = querySum(2*n, tl, mid, st, ed);
            T b = querySum(2*n + 1, mid + 1, tr, st, ed);

            // change me
            return a+b;
        }
    }

    T queryMin(int n, int tl, int tr, int st, int ed) {
        // change me
        if (tl > tr || ed < tl || st > tr) return MAX;
        pushdown(n, tl, tr);

        // segment completely inside query
        if (st <= tl && tr <= ed) {
            return treemin[n];
        }

        int mid = (tl + tr)/2;
        T a = queryMin(2*n, tl, mid, st, ed);
        T b = queryMin(2*n + 1, mid + 1, tr, st, ed);

        // change me
        return min(a, b);
    }
};

const int SEGSZ = 4*MXV;
LazySeg<long long, SEGSZ> seg;

Fast Segment Tree

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
template<class T> struct Seg { // comb(ID,b) = b
    const T ID = 1e18; T comb(T a, T b) { return min(a,b); }
    int n; vector<T> seg;
    void init(int _n) { n = _n; seg.assign(2*n,ID); }
    void pull(int p) { seg[p] = comb(seg[2*p],seg[2*p+1]); }
    void upd(int p, T val) { // set val at position p
        seg[p += n] = val; for (p /= 2; p; p /= 2) pull(p); }
    T query(int l, int r) {	// min on interval [l, r]
        T ra = ID, rb = ID;
        for (l += n, r += n+1; l < r; l /= 2, r /= 2) {
            if (l&1) ra = comb(ra,seg[l++]);
            if (r&1) rb = comb(seg[--r],rb);
        }
        return comb(ra,rb);
    }
};

Seg<int> st;
st.upd(i, x);
int a = st.query(x, y);

LCA

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
int T = 1;
vector<int> st, ed, depth;
vector<vector<int>> up(2e5 + 1, vector<int> (20));

//euler tour to flattern the tree
void dfs(int n, int p) {
    st[n] = T;
    depth[n] = depth[p] + 1;

    //jump table for node n
    up[n][0] = p;
    for (int i = 1; i < 20; i++) {
        up[n][i] = up[up[n][i - 1]][i - 1];
    }

    T++;
    for (auto c : adj[n]) {
        if (c != p) {
            dfs(c, n);
        }
    }

    ed[n] = T - 1;
}

int lca(int a, int b) {
    if (depth[a] < depth[b]) swap(a, b);

    int diff = depth[a] - depth[b];
    for (int i = 0; i < 20; i++) {
        if ((diff >> i) & 1) {
            a = up[a][i];
        }
    }
    if (a == b) return a;

    for (int i = 19; i >= 0; i--) {
        int ap = up[a][i];
        int bp = up[b][i];
        if (ap != bp) {
            a = ap; b = bp;
        }
    }
    return up[a][0];
}

//dfs to get depth and initialize jump table
depth[0] = -1;
dfs(1, 0);

//lca
int l = lca(a, b);

HLD

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
//dfs1 to get depth, parent and sz of each node
void dfs(int n, int p) {
    // make trees
    depth[n] = depth[p] + 1;
    sz[n] = 1;
    for (auto i : adj[n]) {
        if (i != p) {
            parent[i] = n;
            dfs(i, n);
            sz[n] += sz[i];
        }
    }
}

int T = 1;
//dfs2 to flatten and get index T of node, head of node
void dfs2(int n, int h, int p) {
    st[n] = T++;

    //do some calculation.e.g    seg.upd(st[n], values[n]);
    head[n] = h;
    int hv = -1;
    for (auto i : adj[n]) {
        if (i != p) {
            if (hv == -1 || sz[i] > sz[hv]) {
                hv = i;
            }
        }
    }

    // leaf node, then return
    if (hv == -1) return;

    // heavy edge = inherit
    dfs2(hv, h, n);

    // other eddge, start a new segment
    for (auto i : adj[n]) {
        if (i != p && i != hv) {
            // light edge = new segment
            dfs2(i, i, n);
        }
    }
}

//do something for 2 nodes ...
int hldquery(int a, int b) {
    int ans = 0;
    while (head[a] != head[b]) {
        if (depth[head[a]] < depth[head[b]]) swap(a, b);
        int q1 = seg.query(st[head[a]], st[a]);
        a = parent[head[a]];
        ans = max(ans, q1);
    }
    if (depth[a] < depth[b]) swap(a, b);
    int q2 = seg.query(st[b], st[a]);
    ans = max(ans, q2);
    return ans;
}

Union Find

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
int parent[N]; // initialize parent[i] = i

#ifndef PATH_COMPRESSION
int findRoot (int a) {
    if (parent[a] == a) {
      return a;
    }

    return findRoot(parent[a]);
}
#else
//Path Compression - O(log N) per query:
int findRoot (int a) {
  if (parent[a] == a) {
    return a;
  }
    return parent[a] = findRoot(parent[a]);
}
#endif

bool isConnected (int a, int b) {
  return findRoot(a) == findRoot(b);
}

#ifdef KEEP_SIZE_SMALL
int size[N]; // initialize size[i] = 1

void join (int a, int b) {
  a = findRoot(a);
  b = findRoot(b);

  if (a == b) {
    return;
  }

  if (size[a] < size[b]) {
  parent[a] = b;
    size[b] += size[a];
  }
  else {
    parent[b] = a;
    size[a] += size[b];
  }
}
#elseif KEEP_DEPTH_SMALL
int depth[N]; // initialize depth[i] = 1

void join (int a, int b) {
  a = findRoot(a);
  b = findRoot(b);

  if (a == b) {
    return;
  }

  if (depth[a] < depth[b]) {
    parent[a] = b;
  }
  else {
    parent[b] = a;
    depth[a] = max(depth[a], depth[b] + 1);
  }
}
#else
void join (int a, int b) {
  parent[ findRoot(a) ] = findRoot(b);
}

#endif

KMP

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
typedef vector<int> VI;

void buildTable(string& w, VI& t)
{
  t = VI(w.length());
  int i = 2, j = 0;
  t[0] = -1; t[1] = 0;

  while(i < w.length())
  {
    if(w[i-1] == w[j]) {
       t[i] = j+1;
       i++;
       j++;
    }
    else if(j > 0) {
      j = t[j];
    }
    else {
      t[i] = 0;
      i++;
    }
  }
}

int KMP(string& s, string& w)
{
  int m = 0, i = 0;
  VI t;

  buildTable(w, t);
  while(m+i < s.length())
  {
    if(w[i] == s[m+i])
    {
      i++;
      if(i == w.length()) return m;
    }
    else
    {
      m += i-t[i];
      if(i > 0) i = t[i];
    }
  }
  return s.length();
}

Math

Fast Power

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
int mul(int x,int n,int p)
{
   int ans=0;
   while(n) {
       if (n&1)ans=ans+x%p;
       x=x+x%p;
       n>>=1;
   }
   return ans%p;
}

int fpow(int x,int n,int p)
{
   int ans=1;
   while (n) {
       if (n&1) ans=mul(ans,x,p);
       x=mul(x,x,p);
       n>>=1;
   }
   return ans%p;
}

GCD

1
2
3
4
5
6
7
8
9
10
11
12
13
14
int gcd(int x,int y){
    return y==0?x:gcd(y,x%y);
}

int exgcd(int a,int b,int &x,int &y)
{
    if (b==0) {
        x=1;y=0;
        return a;
    }
    int g=exgcd(b,a%b,x,y);
    int t=x;x=y;y=t-a/b*x;
    return g;
}

Data Structure

Unordered Map Custom Comparitor

1
2
3
4
struct hashPi {
    size_t operator()(const pair<int, int>& p) const { return (p.first*100001) + p.second; }
};
unordered_map<pair<int, int>, int, hashPi> findlca;

Priority queue

1
2
3
4
5
6
7
8
9
10
11
12
priority_queue<int> q;                           //biggest on top
priority_queue<int,vector<int>,less<int> > q;    //biggest on top

priority_queue<int,vectot<int>,greater<int> > q; //smallest on top

struct cmp{
    bool operator () (pair<int,int> &a, pair<int,int> &b){
        if (a.first!= b.first) return a.first<b.first;
        else return a.second<b.second;
    }
};
priority_queue<rec,vector<rec>,cmp>q;  

Sorting

Quick Sort

1
2
3
4
5
6
7
8
9
10
11
12
13
void quick_sort(int q[], int l, int r)
{
    if (l >= r) return;

    int i = l - 1, j = r + 1, x = q[l + r >> 1];
    while (i < j)
    {
        do i ++ ; while (q[i] < x);
        do j -- ; while (q[j] > x);
        if (i < j) swap(q[i], q[j]);
    }
    quick_sort(q, l, j), quick_sort(q, j + 1, r);
}

Merge Sort

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
void merge_sort(int q[], int l, int r)
{
    if (l >= r) return;

    int mid = l + r >> 1;
    merge_sort(q, l, mid);
    merge_sort(q, mid + 1, r);

    int k = 0, i = l, j = mid + 1;
    while (i <= mid && j <= r)
        if (q[i] <= q[j]) tmp[k ++ ] = q[i ++ ];
        else tmp[k ++ ] = q[j ++ ];

    while (i <= mid) tmp[k ++ ] = q[i ++ ];
    while (j <= r) tmp[k ++ ] = q[j ++ ];

    for (i = l, j = 0; i <= r; i ++, j ++ ) q[i] = tmp[j];
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
bool check(int x) {/* ... */}  

int bsearch_1(int l, int r)
{
    while (l < r)
    {
        int mid = l + r >> 1;
        if (check(mid)) r = mid;    
        else l = mid + 1;
    }
    return l;
}

int bsearch_2(int l, int r)
{
    while (l < r)
    {
        int mid = l + r + 1 >> 1;
        if (check(mid)) l = mid;
        else r = mid - 1;
    }
    return l;
}

Math

Binary Search for Equation Solving

1
2
3
4
5
6
7
8
9
10
11
12
13
bool check(double x) {/* ... */}  

double bsearch_3(double l, double r)
{
    const double eps = 1e-6;   
    while (r - l > eps)
    {
        double mid = (l + r) / 2;
        if (check(mid)) r = mid;
        else l = mid;
    }
    return l;
}

Big Integer

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
// C = A + B, A >= 0, B >= 0
vector<int> add(vector<int> &A, vector<int> &B)
{
    if (A.size() < B.size()) return add(B, A);

    vector<int> C;
    int t = 0;
    for (int i = 0; i < A.size(); i ++ )
    {
        t += A[i];
        if (i < B.size()) t += B[i];
        C.push_back(t % 10);
        t /= 10;
    }

    if (t) C.push_back(t);
    return C;
}

// C = A - B, A >= B, A >= 0, B >= 0
vector<int> sub(vector<int> &A, vector<int> &B)
{
    vector<int> C;
    for (int i = 0, t = 0; i < A.size(); i ++ )
    {
        t = A[i] - t;
        if (i < B.size()) t -= B[i];
        C.push_back((t + 10) % 10);
        if (t < 0) t = 1;
        else t = 0;
    }

    while (C.size() > 1 && C.back() == 0) C.pop_back();
    return C;
}

// C = A * b, A >= 0, b > 0
vector<int> mul1(vector<int> &A, int b)
{
    vector<int> C;

    int t = 0;
    for (int i = 0; i < A.size() || t; i ++ )
    {
        if (i < A.size()) t += A[i] * b;
        C.push_back(t % 10);
        t /= 10;
    }

    while (C.size() > 1 && C.back() == 0) C.pop_back();

    return C;
}

// C = A * B
vector<int> mul2(vector<int> &A, vector<int> &B) {
    vector<int> C(A.size() + B.size(), 0);  

    for (int i = 0; i < A.size(); i++)
        for (int j = 0; j < B.size(); j++)
            C[i + j] += A[i] * B[j];

    int t = 0;
    for (int i = 0; i < C.size(); i++) {  
        t += C[i];
        C[i] = t % 10;
        t /= 10;
    }

    while (C.size() > 1 && C.back() == 0) C.pop_back();
    return C;
}

// A / b = C ... r, A >= 0, b > 0
vector<int> div(vector<int> &A, int b, int &r)
{
    vector<int> C;
    r = 0;
    for (int i = A.size() - 1; i >= 0; i -- )
    {
        r = r * 10 + A[i];
        C.push_back(r / b);
        r %= b;
    }
    reverse(C.begin(), C.end());
    while (C.size() > 1 && C.back() == 0) C.pop_back();
    return C;
}

Prefix Sum

1D Prefix Sum

1
2
3
4
5
6
for (int i=1; i<=N; i++) {
  pre[i] = val[i] + pre[i-1];
}

//sum from a to b inclusive
pre[b] - pre[a-1]

Reverse 1D Prefix Sum

1
2
3
4
5
6
val[a] += x;
val[b+1] -= x;

for (int i=1; i< N; i++) {
  val[i] += val[i-1];
}

2D Prefix Sum

1
2
3
4
5
6
7
8
for (int i=1; i<=N; i++) {
  for (int j=1; j<=N; j++) {
    pre[i][j] = val[i][j] + pre[i-1][j] + pre[i][j-1] - pre[i-1][j-1];
  }
}

//sum of the range from (x1, y1) to (x2, y2) is:
pre[x2][y2] - pre[x1-1][y2] - pre[x2][y1-1] + pre[x1-1][y1-1]

Reverse 2D Prefix Sum

1
2
3
4
5
6
7
8
9
10
11
12
// add v to the range from (x1, y1) to (x2, y2), inclusive
val[x1][y1] += v;
val[x1][y2+1] -= v;
val[x2+1][y1] -= v;
val[x2+1][y2+1] += v;

// get value at (i,j)
for (int i=1; i<=N; i++) {
  for (int j=1; j<=N; j++) {
    val[i][j] += + pre[i-1][j] + pre[i][j-1] - pre[i-1][j-1];
  }
}

Bit Operation

The k-th bit of an integer n

1
n >> k & 1

The first LSB 1 and it’s trailing 0s

1
lowbit(n) = n & -n

2 Pointers

1
2
3
4
5
6
for (int i = 0, j = 0; i < n; i ++ )
{
    while (j < i && check(i, j)) j ++ ;

    // your code
}

Coordinate Compression

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
vector<int> alls; // store all input values
sort(alls.begin(), alls.end());
alls.erase( unique( alls.begin(), alls.end() ), alls.end()); // remove duplicates

// use binary search to find coordinate of value x
// can also be done using upper_bound(alls.begin(), all.end(), x)
int find(int x)
{
    return upper_bound(alls.begin(), alls.end(), x) - alls.begin();
}

//add is array of pairs, first is the uncompressed index, second is the operation to add
//a is the compressed array, each index is compressed from original index
map<int, int> a;
for(auto item : add)
{
    int compressed_index = find(item.first);
    a[compressed_index] += item.second;
}

Segments Merger

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
// combine the over lapping ranges
void merge(vector<pair<int,int>> &segs)
{
    vector<pair<int,int>> res;

    sort(segs.begin(), segs.end());

    int st = -2e9, ed = -2e9;

    //[st, ed]: the last range
    for (auto seg : segs)
        //case 1: can't merge new range
        if (ed < seg.first)
        {
            if (st != -2e9)
              res.push_back({st, ed});

            st = seg.first, ed = seg.second;
        }
        else
          //case 2: the new range's end is smaller than last range's end
          //case 3: the new range's end is bigger than last range's end
          //in both cases, extend the new end for the last range
          ed = max(ed, seg.second);

    if (st != -2e9)
      res.push_back({st, ed});

    segs = res;
}

Monotonous Stack

1
2
3
4
5
6
7
8
9
10
11
12
//Open a stack and decrease monotonically.
//For the top element on the stack, what you can see is the number of elements behind
	stack<int>s;
	for (int i = 0; i < n; i++) {
		int temp;
		cin >> temp;

		//When the input is bigger, remove the smaller from the stack
		while (!s.empty() && temp >= s.top())s.pop();

		s.push(temp);
	}

Sliding Window

Use monotonous queue to find minimal in a sliding window.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
deque<int> dq;

//find the minimal number in the sliding window
for(int i = 0; i < n; i ++)
{
     //keep the front of the deque always inside sliding window of size k
     if( !dq.empty() && k < i - dq.front() + 1)
         dq.pop_front();

     //maintain monotonous queue
     //if the queue is not empty and the tail element is not smaller than a[i]
     //pop all the bigger elements until it is smaller than a[i]
     while( !dq.empty() && a[dq.back()] >= a[i])
         dq.pop_back();

     dq.push_back(i);

     if(i + 1 >= k)
         cout << a[dq.front()] << " ";
}

KMP

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
//ne: based on pattern string p
//ne[i] = j, j the maximal length of prefix: ne[1:j] == ne[i-j+1:i]
for (int i = 2, j = 0; i <= m; i ++ )
{
    while (j && p[i] != p[j + 1])
      j = ne[j];

    if (p[i] == p[j + 1])
      j ++ ;

    ne[i] = j;
}

for (int i = 1, j = 0; i <= n; i ++ )
{
    while (j && s[i] != p[j + 1])
      j = ne[j];

    if (s[i] == p[j + 1])
      j++;

    if (j == m)
    {
        j = ne[j];
    }
}

Trie

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
int son[N][26], cnt[N], idx;
// 0 is the root, and empty node  
// son[p][u] stores p's subtree root node
// cnt[p] stores number of words ending with p node

void insert(string str)
{
    int p = 0;

    for (int i = 0; i < str.size(); i ++ )
    {
        int u = str[i] - 'a';

        if (!son[p][u])
          son[p][u] = ++idx;

        p = son[p][u];
    }

    cnt[p]++ ;
}

int query(string str)
{
    int p = 0;
    for (int i = 0; str.size(); i ++ )
    {
        int u = str[i] - 'a';

        if (!son[p][u])
          return 0;

        p = son[p][u];
    }

    return cnt[p];
}
Back to top ↑