题解 | #接头密匙#

#include <array>
#include <cctype>
#include <string>
#include <vector>
class Solution {
public:
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     * 
     * @param b int整型vector<vector<>> 
     * @param a int整型vector<vector<>> 
     * @return int整型vector
     */
    vector<int> countConsistentKeys(vector<vector<int> >& b, vector<vector<int> >& a) {
        // write code here
        // a[0] 的长度比较大
        std::string str;
        int diff;
        std::vector<int> res;
        for (auto &data : a) {
            str.clear();
            for (int i = 1; i < data.size(); i++) {
                diff = data[i] - data[i - 1];
                str.append(std::to_string(diff));
                str.push_back('#');
            }
            tree.insert(str);
        }

        for (auto &data : b) {
            str.clear();
            for (int i = 1; i < data.size(); i++) {
                diff = data[i] - data[i - 1];
                str.append(std::to_string(diff));
                str.push_back('#');
            }
            res.emplace_back(tree.prefixCounts(str));
        }
        tree.clear();
        return res;
    }

private:
    // 添加一个前缀树的类
    class Tree
    {
        static constexpr int MAXN = 2*(1e6) + 1;  //  填1e7内存直接爆了
        std::array<std::array<int, 12>, MAXN> nodes; // 12个字符 [0, 9], -, #
        std::array<int, MAXN> pass;
        std::array<int, MAXN> end;
        int cnt = 1;
    public:
        Tree()
        {
            clear();
        }

        void clear()
        {
            cnt = 1;
            for (auto &item : nodes) {
                item.fill(0);
            }
            pass.fill(0);
            end.fill(0);
        }

        int getPath(char ch)
        {
            if (std::isdigit(ch)) {
                return ch - '0';
            }

            if (ch == '-') {
                return 10;
            }

            // '#'
            return 11;
        }

        void insert(const std::string &s)
        {
            int cur = 1;
            pass[cur]++;
            for (char ch : s) {
                int path = getPath(ch);
                if (nodes[cur][path] == 0) {
                    nodes[cur][path] = ++cnt;
                }
                cur = nodes[cur][path];
                pass[cur]++;
            }
            end[cur]++;
        }

        // 暂时用不上
        // void remove(const std::string &s)
        // {

        // }

        // void counts(const std::string &s)
        // {

        // }

        int prefixCounts(const std::string &s)
        {
            int cur = 1;
            int path;
            for (char ch : s)
            {
                path = getPath(ch);
                if (nodes[cur][path] == 0) {
                    return 0;
                }
                cur = nodes[cur][path];
            }
            return pass[cur];
        }
    };

    Tree tree;
};