题解 | 接头密匙

#include<bits/stdc++.h>
using namespace std;
const int MAXN = 2e6 + 10;

class Solution {
  public:
    int trieTree[MAXN][12];
    int pass[MAXN];
    int cnt;
    void build() {
        cnt = 1;
    }

    int getPath(char c) {
        if (c == '-')return 10;
        else if (c == '#')return 11;
        else {
            return c - '0';
        }
    }

    void insert(const string& word) {
        int cur = 1;
        pass[cur]++;
        for (char c : word) {
            int path = getPath(c);
            if (trieTree[cur][path] == 0) {
                trieTree[cur][path] = ++cnt;
            }
            cur = trieTree[cur][path];
            pass[cur]++;
        }
    }

    int prefixNumber(const string& prefix) {
        int cur = 1;
        for (char c : prefix) {
            int path = getPath(c);
            if (trieTree[cur][path] == 0) {
                return 0;
            }
            cur = trieTree[cur][path];
        }
        return pass[cur];
    }

    void clear() {
        for (int i = 1; i <= cnt; i++) {
            for (int j = 0; j < 12; j++) {
                trieTree[i][j] = 0;
            }
            pass[i] = 0;
        }
    }
    vector<int> countConsistentKeys(vector<vector<int> >& b,
                                    vector<vector<int> >& a) {
        build();
        vector<int>ans;
        for (int i = 0; i < a.size(); i++) {
            string str;
            for (int j = 1; j < a[i].size(); j++) {
                str += to_string(a[i][j] - a[i][j - 1]) + "#";
            }
            insert(str);
        }
        for (int i = 0; i < b.size(); i++) {
            string str;
            for (int j = 1; j < b[i].size(); j++) {
                str += to_string(b[i][j] - b[i][j - 1]) + "#";
            }
            ans.push_back(prefixNumber(str));

        }
        clear();
        return ans;
    }
};