算法练习常用模板

字符串

回文子串

动态规划

1
2
3
4
5
6
7
8
9
10
11
12
// java version
// yes[i][j] 表示下标从 i 到 j 的子串是否为回文子串
int len = s.length();
boolean[][] yes = new boolean[len][len];
for (int i = 0; i < len; i++) {
    Arrays.fill(yes[i], true);
}
for (int i = len-1; i >= 0; i--) {
    for (int j = i+1; j < len; j++) {
        yes[i][j] = (s.charAt(i) == s.charAt(j) && yes[i+1][j-1]);
    }
}

Manacher 算法

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
class Solution {
public:
    int countSubstrings(string s) {
        string now = "$#";
        for (auto ch : s) {
            now += ch;
            now += '#';
        }
        vector<int> f(now.size(), 0);
        int mi = 0, mr = 0, ans = 0, len = now.size();
        now += '!';
        for (int i = 1; i < len; i++) {
            f[i] = (i <= mr ? min(f[2*mi - i], mr - i + 1):1);
            while (now[i + f[i]] == now[i - f[i]]) {
                f[i]++;
            }
            if (i + f[i] - 1 > mr) {
                mi = i;
                mr = i + f[i] - 1;
            }
            ans += (f[i] - 1)%2 ? f[i]/2:(f[i] - 1)/2;
        }
        return ans;
    }
};

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
vector<int> prefix_function(string t) {
    int n = t.length();
    vector<int> next(n);
    for (int i = 1; i < n; i++) {
        int j = next[i - 1];
        while (j > 0 && t[i] != t[j]) {
            j = next[j - 1];
        }
        if (t[i] == t[j]) {
            j++;
        }
        next[i] = j;
    }
    return next;
}

vector<pair<int, int>> KMP(string s, string t) {
    int i = 0, j = 0;
    vector<int> next = prefix_function(t);
    vector<pair<int, int>> ans;
    while (i < s.length()) {
        if (s[i] == t[j]) {
            i++, j++;
            if (j == t.length()) {
                ans.emplace_back(make_pair(i - j, i - 1));
                j = next[j - 1];
            }
        }
        else {
            if (j > 0) {
                j = next[j - 1];
            }
            else {
                i++;
            }
        }
    }
    return ans;
}

Z 函数 - 扩展 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
/**
    z[i] 表示 s[i:] 与 s 的最长公共前缀长度
 */
vector<int> zFunc(string s) {
    int n = s.size();
    vector<int> z(n, 0);
    int l = 0, r = 0;
    for (int i = 1; i < n; i++) {
        if (i <= r && z[i - l] < r - i + 1) {
            z[i] = z[i - l];
        } else {
            z[i] = max(0, r - i + 1);
            while (i + z[i] < n && s[z[i]] == s[i + z[i]]) {
                z[i]++;
            }
        }
        if (i + z[i] - 1 > r) {
            l = i;
            r = i + z[i] - 1;
        }
    }
    return z;
}

// version 2:
vector<int> zFunc(string s) {
    int n = s.size();
    vector<int> z(n, 0);
    int l = 0, r = 0;
    for (int i = 1; i < n; i++) {
        if (i <= r) {
            z[i] = min(z[i - l], r - i + 1);
        }
        while (i + z[i] < n && s[z[i]] == s[i + z[i]]) {
            l = i;
            r = i + z[i];
            z[i]++;
        }
    }
    return z;
}

AC 自动机

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
class ACTree {

private:
    ACTree *child[26];
    ACTree *fail;
    vector<int> exist; // exist 数组存储的是以当前结尾的单词长度

public:
    ACTree() {
        fail = nullptr;
        for (int i = 0; i < 26; i++) {
            child[i] = nullptr;
        }
    }

    void insert(const string &s) {
        ACTree *cur = this;
        for (auto &ch : s) {
            int idx = ch - 'a';
            if (cur->child[idx] == nullptr) {
                cur->child[idx] = new ACTree();
            }
            cur = cur->child[idx];
        }
        cur->exist.push_back(s.length());
    }

    void build(const vector<string> &s) {
        int n = s.size();
        for (auto &word : s) {
            insert(word);
        }
        queue<ACTree*> q;
        for (int i = 0; i < 26; i++) {
            if (this->child[i] != nullptr) {
                this->child[i]->fail = this;
                q.push(this->child[i]);
            }
        }
        while (!q.empty()) {
            ACTree *cur = q.front();
            q.pop();
            for (int i = 0; i < 26; i++) {
                if (cur->child[i] != nullptr) {
                    ACTree *nxt = cur->child[i], *fafail = cur->fail;
                    while (fafail != nullptr && fafail->child[i] == nullptr) {
                        fafail = fafail->fail;
                    }
                    if (fafail == nullptr) {
                        nxt->fail = this;
                    }
                    else {
                        nxt->fail = fafail->child[i];
                    }
                    if (nxt->fail->exist.size() != 0) {
                        for (auto &it : nxt->fail->exist) {
                            nxt->exist.push_back(it);
                        }
                    }
                    q.push(nxt);
                }
            }
        }
    }

    // 返回的是所有在 t 中出现的单词的左右下标
    vector<pair<int, int>> query(const string &t) {
        ACTree *cur = this;
        vector<pair<int, int>> ans;
        for (int i = 0; i < t.length(); i++) {
            int idx = t[i] - 'a';
            while (cur->child[idx] == nullptr && cur->fail != nullptr) {
                cur = cur->fail;
            }
            if (cur->child[idx] != nullptr) {
                cur = cur->child[idx];
            }
            else {
                continue;
            }
            if (cur->exist.size() != 0) {
                for (auto &it : cur->exist) {
                    ans.emplace_back(make_pair(i - it + 1, i));
                }
            }
        }
        return ans;
    }
};

括号匹配

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
// 处理括号字符串 s 的通用方法：
// 遍历 s，用一个变量 cnt 记录 s 前缀的左括号的个数减去右括号的个数。
// 如果 s 是合法括号字符串，那么在遍历的过程中 cnt >= 0，且遍历结束后 cnt = 0。

// 把拼接的字符串分别记作 S 和 T，即 S + T 是合法括号字符串。
// 设 S 的 cnt 值为 cntS，要求遍历过程中的 cntS 的最小值 >= 0。
// 设 T 的 cnt 值为 cntT，要求遍历过程中的 cntT 的最小值等于遍历结束后的 cntT。
// 如果 S+T 是合法括号字符串，那么 cntS + cntT = 0。

// Code:
// mp 存储的是成为合法括号字符串对应的冗余量，
// 正值表示该字符串应该位于左边，负值为右边，零值表示该字符串为合法括号字符串
map<int, int> mp;
for (int i = 0; i < n; i++) {
    string s;
    cin >> s;
    int len = s.length(), cnt = 0, mcnt = len;
    for (auto ch : s) {
        cnt += (ch == '(' ? 1 : -1);
        mcnt = min(cnt, mcnt);
    }
    if (cnt >= 0 && mcnt >= 0 || cnt < 0 && mcnt == cnt) {
        mp[cnt]++;
    }
}

字符串哈希

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
class stringHash {
private:
    unsigned long long *hash;
    vector<int> src;
    int length;
    int p;
    unsigned long long *exp_p;

public:
    stringHash(vector<int> &src, int p) {
        length = src.size();
        this->src = src;
        hash = new unsigned long long[length];
        exp_p = new unsigned long long[length];
        this->p = p;
        // 初始化前缀hash 和 p^i
        hash[0] = src[0];
        exp_p[0] = 1;
        for (int i = 1; i < length; i++) {
            hash[i] = hash[i-1]*p + src[i];
            exp_p[i] = p*exp_p[i-1];
        }
    }

    // 0 <= l <= r <= length
    unsigned long long query(int l, int r) {
        return (l == 0 ? hash[r]:hash[r] - hash[l-1]*exp_p[r-l+1]);
    }

    bool isSame(unsigned long long src, int l, int r) {
        return src == query(l, r);
    }
};

树

前序 + 中序建树

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution {
public:
    TreeNode* buildTree(vector<int>& preorder, vector<int>& inorder) {
        unordered_map<int, int> idx;
        for (int i = 0; i < inorder.size(); i++) {
            idx[inorder[i]] = i;
        }
        function<TreeNode*(int, int, int, int)> dfs = [&](int pl, int pr, int il, int ir) -> TreeNode* {
            if (pl > pr || il > ir) {
                return nullptr;
            }
            TreeNode *cur = new TreeNode(preorder[pl]);
            int m = idx[preorder[pl]];
            int len = m - il;
            TreeNode *l = dfs(pl + 1, pl + len, il, il + len - 1);
            TreeNode *r = dfs(pl + len + 1, pr, m + 1, ir);
            cur->left = l, cur->right = r;
            return cur;
        };
        return dfs(0, preorder.size() - 1, 0, preorder.size() - 1);
    }
};

后序 + 中序建树

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
public:
    TreeNode* buildTree(vector<int>& inorder, vector<int>& postorder) {
        unordered_map<int, int> pos;
        for (int i = 0; i < inorder.size(); i++) {
            pos[inorder[i]] = i;
        }
        function<TreeNode*(int, int, int, int)> dfs = [&](int il, int ir, int pl, int pr) -> TreeNode* {
            if (il > ir || pl > pr) {
                return nullptr;
            }
            int m = pos[postorder[pr]];
            int len = m - il;
            TreeNode *cur = new TreeNode(postorder[pr]);
            cur->left = dfs(il, m - 1, pl, pl + len - 1);
            cur->right = dfs(m + 1, ir, pl + len, pr - 1);
            return cur;
        };
        return dfs(0, inorder.size() - 1, 0, inorder.size() - 1);
    }
};

前序 + 后序建树

根据数据结构知识，树并不唯一，返回其中一种情况

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
public:
    TreeNode* constructFromPrePost(vector<int>& preorder, vector<int>& postorder) {
        unordered_map<int, int> idx;
        for (int i = 0; i < postorder.size(); i++) {
            idx[postorder[i]] = i;
        }
        function<TreeNode*(int, int, int, int)> solve = [&](int lpre, int rpre, int lpost, int rpost) -> TreeNode* {
            if (lpre > rpre || lpost > rpost) {
                return nullptr;
            }
            if (lpre == rpre) {
                return new TreeNode(preorder[lpre]);
            }
            int m = idx[preorder[lpre + 1]];
            int len = m - lpost;
            TreeNode *cur = new TreeNode(preorder[lpre]);
            cur->left = solve(lpre + 1, lpre + 1 + len, lpost, m);
            cur->right = solve(lpre + 1 + len + 1, rpre, m + 1, rpost - 1);
            return cur;
        };
        return solve(0, preorder.size() - 1, 0, preorder.size() - 1);
    }
};

树状数组

注释部分为区间最值

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
class Fenwick {
    vector<long long> tree;
    int n;
public:
    // max -> tree(n + 1, LLONG_MIN)
    Fenwick(int n) : tree(n + 1, 0), n(n) {}

    inline int lowbit(int x) {
        return x&-x;
    }

    // max -> update(int, long long)
    void add(int x, long long val) {
        for (int i = x; i <= n; i += lowbit(i)) {
            // tree[i] = max(tree[i], val);
            tree[i] += val;
        }
    }

    long long pre_sum(int x) {
        long long res = 0;
        for (int i = x; i > 0; i -= lowbit(i)) {
            res += tree[i];
        }
        return res;
    }

    // long long pre_max(int i) {
    //     long long res = LLONG_MIN;
    //     while (i > 0) {
    //         res = max(res, tree[i]);
    //         i -= i & -1;
    //     }
    //     return res;
    // }
};

图

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
34
35
/**
 * @brief 计算单源最短路径以及最短路径数量
 * 
 * @param graph 
 * @param start 
 * @param mod 
 * @return {最短路径长度, 最短路径数量}
 */
pair<vector<long long>, vector<int>> dijkstra(vector<vector<pair<int, int>>> &graph, int start, const int mod) {
    int n = graph.size();
    vector<long long> dist(n, LLONG_MAX);
    vector<int> cnts(n, 0);
    dist[start] = 0;
    cnts[start] = 1;
    priority_queue<pair<long long, int>, vector<pair<long long, int>>, greater<>> pq;
    pq.push({dist[start], start});
    while (!pq.empty()) {
        auto [d, x] = pq.top();
        pq.pop();
        if (d > dist[x]) {
            continue;
        }
        for (auto [to, val] : graph[x]) {
            long long cur_d = dist[x] + val;
            if (cur_d < dist[to]) {
                dist[to] = cur_d;
                cnts[to] = cnts[x];
                pq.push({cur_d, to});
            } else if (cur_d == dist[to]) {
                cnts[to] = (cnts[to] + cnts[x])%mod;
            }
        }
    }
    return {dist, cnts};
}

数学

素数筛

1
2
3
4
5
6
7
8
9
10
11
12
13
14
// 埃式筛法
const int MX = 1e5;
vector<bool> isPrime(MX + 5, true);
int init = []() {
    isPrime[0] = isPrime[1] = false;
    for (int i = 2; i * i <= MX; i++) {
        if (isPrime[i]) {
            for (int j = i * i; j <= MX; j += i) {
                isPrime[j] = false;
            }
        }
    }
    return 0;
}();

质因数分解

需要预处理质数集合 prime

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
vector<pair<int, int>> prime_factorization(int a) {
    vector<pair<int, int>> fact;
    for (int j = 0; j < prime.size() && prime[j] < cur; j++) {
        int cnt = 0;
        if (cur % prime[j] == 0) {
            while (cur % prime[j] == 0) {
                cur /= prime[j];
                cnt++;
            }
            fact.push_back({prime[j], cnt});
        }
    }
    if (cur != 1) {
        fact.push_back({cur, 1});
    }
}

矩阵快速幂

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
const int mod = 1e9 + 7;

// calc A*B
vector<vector<long long int>> multiply(vector<vector<long long int>> &a, vector<vector<long long int>> &b) {
    int m = a.size(), n = b[0].size();
    if (a[0].size() != b.size()) {
        return {{}};
    }
    int t = a[0].size();
    vector<vector<long long int>> ans(m, vector<long long int>(n));
    for (int i = 0; i < m; i++) {
        for (int j = 0; j < n; j++) {
            ans[i][j] = 0;
            for (int k = 0; k < t; k++) {
                ans[i][j] = (ans[i][j] + a[i][k]*b[k][j])%mod;
            }
        }
    }
    return ans;
}

// calc A^k
vector<vector<long long int>> qpow(vector<vector<long long int>> &a, long long int k) {
    if (a.size() != a[0].size()) {
        return {{}};
    }
    int n = a.size();
    vector<vector<long long int>> ans(n, vector<long long int>(n, 0));
    for (int i = 0; i < n; i++) {
        ans[i][i] = 1;
    }
    while (k) {
        if (k & 1) {
            ans = multiply(ans, a);
        }
        k >>= 1;
        a = multiply(a, a);
    }
    return ans;
}

数据处理

前缀和

设数组 \(\operatorname{nums}\) 的前缀和为 \(\operatorname{pre\_sum}\)，其中

\[ \operatorname{pre\_sum}[i]=\sum_{j=0}^{i-1} \operatorname{nums}[j] \]

因此区间 \([L, R]\) 的前缀和表示为 \(\operatorname{pre\_sum}[R + 1] - \operatorname{pre\_sum}[L]\)

1
2
3
4
5
6
7
8
vector<int> pre_sum(n + 1, 0);
for (int i = 1; i <= n; i++) {
    pre_sum[i] = pre_sum[i - 1] + nums[i - 1];
}

auto calc = [&](int L, int R) {
    return pre_sum[R + 1] - pre_sum[L];
}

前缀和的前缀和

设数组 \(\operatorname{nums}\) 的前缀和的前缀和为 \(\operatorname{pre\_sum\_sum}\)，其中

\[ \operatorname{pre\_sum\_sum}[i]=\sum_{j=0}^{i-1} \operatorname{pre\_sum}[j] \]

因此区间 \([L, R]\) 的前缀和的前缀和表示为 \(\operatorname{pre\_sum\_sum}[R + 1] - \operatorname{pre\_sum\_sum}[L]\)

1
2
3
4
5
6
7
8
9
10
vector<int> pre_pre_sum(n + 2, 0);
long long int pre_sum = 0;
for (int i = 1; i <= n; i++) {
    pre_sum += nums[i - 1];
    pre_pre_sum[i + 1] = pre_pre_sum[i] + pre_sum;
}

auto calc = [&](int L, int R) {
    return pre_pre_sum[R + 1] - pre_pre_sum[L];
}

应用

Problem: (leetcode) 2281. 巫师的总力量和

在范围内含下标 i 的所有子数组的元素和的和可以表示为

\[ \begin{aligned} & \sum_{r=i+1}^{R+1} \sum_{l=L}^i(\operatorname{pre\_sum}[r]-\operatorname{pre\_sum}[l]) \\ = & \left(\sum_{r=i+1}^{R+1} \sum_{l=L}^i \operatorname{pre\_sum}[r]\right)-\left(\sum_{r=i+1}^{R+1} \sum_{l=L}^i \operatorname{pre\_sum}[l]\right) \\ = & (i-L+1) \cdot \sum_{r=i+1}^{R+1} \operatorname{pre\_sum}[r]-(R-i+1) \cdot \sum_{l=L}^i \operatorname{pre\_sum}[l] \end{aligned} \]

因此我们还需要计算出前缀和 \(\operatorname{pre\_sum}\) 的前缀和 \(\operatorname{pre\_sum\_sum}\) ，其中

\[ \operatorname{pre\_sum\_sum}[i]=\sum_{j=0}^{i-1} \operatorname{pre\_sum}[j] \]

综上，区间 [L, R] 中含下标 i 的所有子数组的元素和的和可以表示为:

\[ (i-L+1) \cdot (\operatorname{pre\_sum\_sum}[R+2] - \operatorname{pre\_sum\_sum}[i+1])-(R-i+1) \cdot (\operatorname{pre\_sum\_sum}[i+1]-\operatorname{pre\_sum\_sum}[L]) \]

离散化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
// version 1: arr -> rank(arr)
vector<int> arrayRankTransform(vector<int>& arr) {
    vector<int> tmp = arr;
    sort(tmp.begin(), tmp.end());
    // 去重
    tmp.erase(unique(tmp.begin(), tmp.end()), tmp.end());
    for (auto &x : arr) {
        x = upper_bound(tmp.begin(), tmp.end(), x) - tmp.begin();
    }
    return arr;
}

// version 2: idx dict
unordered_map<int, int> discretization(vector<int>& nums) {
    int n = nums.size();
    auto tmp = nums;
    sort(tmp.begin(), tmp.end());
    tmp.erase(unique(tmp.begin(), tmp.end()), tmp.end());
    unordered_map<int, int> idx;
    for (int i = 0; i < tmp.size(); i++) {
        idx[tmp[i]] = i + 1;
    }
    return idx;
}

数位 DP

视频讲解，从 19:30 开始 https://www.bilibili.com/video/BV1rS4y1s721

其他资料：

https://zhuanlan.zhihu.com/p/348851463

https://www.bilibili.com/video/BV1MT4y1376C

https://www.bilibili.com/video/BV1yT4y1u7jW

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
class Solution {
public:
    int countSpecialNumbers(int n) {
        auto s = to_string(n);
        int m = s.length(), dp[m][1 << 10];
        memset(dp, -1, sizeof(dp));
        // mask 表示当前已经使用的数字，is_num 表示是否填充数字，is_limit 表示当前数位是否受上一数位限制
        // is_num 用于限制前导零
        function<int(int, int, bool, bool)> f = [&](int i, int mask, bool is_limit, bool is_num) -> int {
            if (i == m) {
                return is_num;
            }
            if (!is_limit && is_num && dp[i][mask] >= 0) {
                return dp[i][mask];
            }
            int res = 0;
            // 可以跳过当前数位
            if (!is_num) {
                res = f(i + 1, mask, false, false); 
            }
            // 枚举要填入的数字 d
            int up = is_limit ? s[i] - '0' : 9;
            for (int d = 1 - is_num; d <= up; ++d) {
                // d 不在 mask 中 (非必要条件)
                if ((mask >> d & 1) == 0) {
                    res += f(i + 1, mask | (1 << d), is_limit && d == up, true);
                }
            }
            if (!is_limit && is_num) {
                dp[i][mask] = res;
            }
            return res;
        };
        return f(0, 0, true, false);
    }
};

位运算

HighBit

返回该数对应二进制的最高位 1 所对应的数

1
2
3
4
5
6
7
8
int highbit(int i) {
    i |= (i >>  1);
    i |= (i >>  2);
    i |= (i >>  4);
    i |= (i >>  8);
    i |= (i >> 16);
    return i - (i >>> 1);
}

LowBit

返回该数对应二进制的最低位 1 所对应的数

1
2
3
int lowbit(int i) {
    return i & -i;
}

二进制移除最低位

每次循环将数字 k 的最低二进制为 1 的比特位置为 0

1
2
3
for (; k; k &= k - 1) {
    // code...
}

二进制子集遍历

用于遍历数字 k 的所有二进制子集

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
// version 1:
int sub = k;
do {
    // code...
    sub = (sub - 1) & k;
} while(sub != k);

// version 2:
int sub = k;
while (sub > 0) {
    // code...
    sub = (sub - 1) & k;
}

// version 3:
for (int sub = s; sub; sub = (sub - 1) & s) {
    // code...
}

C++ 位运算内置函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
// 返回 x 的二进制表示中的 1 的个数
int __builtin_popcount(unsigned int x);
int __builtin_popcountl(unsigned long x);
int __builtin_popcountll(unsigned long long x);
// 返回 x 的二进制表示中的 1 的个数奇偶性，偶数个 1 返回 0，反之 1
int __builtin_parity(unsigned int x);
int __builtin_parityl(unsigned long x);
int __builtin_parityll(unsigned long long x);
// 返回 x 的二进制表示中前导的 0 的个数，常用 32 - __builtin_clz(x) 计算 x 对应二进制的有效位数
int __builtin_clz(unsigned int x);
int __builtin_clzl(unsigned long x);
int __builtin_clzll(unsigned long long x);
// 返回 x 的二进制表示中末尾的 0 的个数，不常用
int __builtin_ctz(unsigned int x);
int __builtin_ctzl(unsigned long x);
int __builtin_ctzll(unsigned long long x);

倍增

树上倍增

Problem: (leetcode) 1483. 树节点的第 K 个祖先

通过对 k 的二进制转换能够在 log(k) 的时间复杂度内计算出节点的 k 祖先

预处理时间复杂度：O(nlogn)，询问时间复杂度：O(logk)

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
class TreeAncestor {
private:
    // pa[i][j] 表示节点 i 的第 2^j 个祖先
    vector<vector<int>> pa;
public:
    TreeAncestor(int n, vector<int>& parent) {
        int m = 32 - __builtin_clz(n);
        pa.resize(n, vector<int>(m, -1));
        // 预处理每个节点的第 1 个祖先
        for (int i = 0; i < n; i++) {
            pa[i][0] = parent[i];
        }
        // pa[x][i + 1] = pa[pa[x][i]][i]
        // 即表示 x 的第 2^i 个祖先节点的第 2^i 个祖先节点就是 x 的第 2^(i+1) 个祖先节点。
        for (int i = 0; i < m - 1; i++) {
            for (int j = 0; j < n; j++) {
                int p = pa[j][i];
                if (p != -1) {
                    pa[j][i + 1] = pa[p][i];
                }
            }
        }
    }
    
    int getKthAncestor(int node, int k) {
        int m = 32 - __builtin_clz(k);
        // 二进制分解 k 进行 log(k) 次跳跃查询
        for (int i = 0; i < m; i++) {
            if ((k >> i) & 1) {
                node = pa[node][i];
                if (node < 0) {
                    break;
                }
            }
        }
        return node;
    }
};

【LCA】最近公共祖先

朴素的最近公共祖先可以通过 DFS 进行查询，时间复杂度为 \(O(n)\)。

基于倍增的解法是首先预处理所有节点的第 \(2^i\) 个祖先，然后将待求解的节点 x 和 y 置于同一层，进行 \(log_{2}{k}\) 次祖先查询，最终得到公共祖先，其中 k 为树的最大高度。

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
class TreeAncestor {
    vector<int> depth;
    vector<vector<int>> pa;
public:
    TreeAncestor(vector<pair<int, int>> &edges) {
        int n = edges.size() + 1;
        // n 的二进制长度
        int m = 32 - __builtin_clz(n); 
        vector<vector<int>> g(n);
        // 建图，节点编号从 0 开始
        for (auto [x, y]: edges) {
            g[x].push_back(y);
            g[y].push_back(x);
        }

        depth.resize(n);
        pa.resize(n, vector<int>(m, -1));
        // O(n) 求每个节点的深度，用于后续求 LCA 前转移两个节点为同一层
        function<void(int, int)> dfs = [&](int x, int fa) {
            pa[x][0] = fa;
            for (int y: g[x]) {
                if (y != fa) {
                    depth[y] = depth[x] + 1;
                    dfs(y, x);
                }
            }
        };
        dfs(0, -1);

        // 倍增预处理出每个节点的第 2^i 个祖先
        for (int i = 0; i < m - 1; i++) {
            for (int x = 0; x < n; x++) {
                int p = pa[x][i];
                if (p != -1) {
                    pa[x][i + 1] = pa[p][i];
                }
            }
        }
    }

    // 获取节点 node 的第 k 个祖先
    int get_kth_ancestor(int node, int k) {
        for (; k; k &= k - 1)
            node = pa[node][__builtin_ctz(k)];
        return node;
    }

    // 返回 x 和 y 的最近公共祖先（节点编号从 0 开始）
    int get_lca(int x, int y) {
        if (depth[x] > depth[y])
            swap(x, y);
        // 使 y 和 x 在同一深度
        y = get_kth_ancestor(y, depth[y] - depth[x]);
        if (y == x)
            return x;
        // 循环从大到小，尝试每次上跳 i 步，理解为二进制分解两个节点的公共祖先高度
        for (int i = pa[x].size() - 1; i >= 0; i--) {
            int px = pa[x][i], py = pa[y][i];
            // 跳跃
            if (px != py) {
                x = px;
                y = py;
            }
        }
        // 保证两个节点最终跳跃的位置对应的父节点为最近公共祖先
        return pa[x][0];
    }
};

区间询问

区间数字出现频率

Problem: (leetcode) 2080. 区间内查询数字的频率

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class RangeFreqQuery {
private:
    unordered_map<int, vector<int>> um;

public:
    RangeFreqQuery(vector<int>& arr) {
        for (int i = 0; i < arr.size(); i++) {
            um[arr[i]].emplace_back(i);
        }
    }
    
    int query(int left, int right, int value) {
        if (um.find(value) == um.end()) {
            return 0;
        }
        auto l = lower_bound(um[value].begin(), um[value].end(), left);
        auto r = upper_bound(um[value].begin(), um[value].end(), right);
        return r - l;
    }
};

数据结构

集合前 k 小

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
// 两个 multiset 维护集合中的前 K 小值
struct KNums {
    // 当前维护集合的前 K 小
    int K;
    // st1 保存前 K 小值，st2 保存其它值
    multiset<long long> st1, st2;
    // sum 表示 st1 中所有数的和
    long long sum;

    KNums(int K): K(K), sum(0) {}

    // 调整 st1 和 st2 的大小，保证调整后 st1 保存前 K 小值
    void adjust() {
        while (st1.size() < K && st2.size() > 0) {
            long long t = *(st2.begin());
            st1.insert(t);
            sum += t;
            st2.erase(st2.begin());
        }
        while (st1.size() > K) {
            long long t = *prev(st1.end());
            st2.insert(t);
            st1.erase(prev(st1.end()));
            sum -= t;
        }
    }

    // 插入元素 x
    void add(long long x) {
        if (!st2.empty() && x >= *(st2.begin())) {
            st2.insert(x);
        } else {
            st1.insert(x);
            sum += x;
        }
        adjust();
    }

    // 删除元素 x
    void del(long long x) {
        auto it = st1.find(x);
        if (it != st1.end()) {
            st1.erase(it);
            sum -= x;
        } else {
            st2.erase(st2.find(x));
        }
        adjust();
    }
};

数据流中的中位数

不支持删除操作，支持删除操作的可以转换为『集合前 K 小问题』

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
class MedianFinder {
public:
    priority_queue<int, vector<int>, less<int>> queMin;
    priority_queue<int, vector<int>, greater<int>> queMax;

    MedianFinder() {}

    void addNum(int num) {
        if (queMin.empty() || num <= queMin.top()) {
            queMin.push(num);
            if (queMax.size() + 1 < queMin.size()) {
                queMax.push(queMin.top());
                queMin.pop();
            }
        } else {
            queMax.push(num);
            if (queMax.size() > queMin.size()) {
                queMin.push(queMax.top());
                queMax.pop();
            }
        }
    }

    double findMedian() {
        if (queMin.size() > queMax.size()) {
            return queMin.top();
        }
        return (queMin.top() + queMax.top()) / 2.0;
    }
};
Algorithm
#Algorithm
算法练习常用模板
http://shijieq.github.io/2023/01/01/Algorithm/算法练习常用模板/
Author
ShijieQ
Posted on
January 1, 2023
Licensed under