接頭密匙

推薦事前閱讀：前綴樹/字典樹介紹

思路

題目給定 $n$ 個鑰匙屬於 $A$ ， $m$ 個鑰匙屬於 $B$ ，
配對的滿足條件是 $b[i+1]-b[i] = a[i+1]-a[i]$ ，其中 $a\in{A}, b\in{B}$ ，鑰匙 $b$ 的長度比 $a$ 小。
問每一把鑰匙 $b$ 能配對到多少 $A$ 中的鑰匙。
就是問鑰匙 $\gray{b}$ 的前綴跟 $\gray{a}$ 有多少

比如b = [1,2,3,4,5]，a = [3,4,5,6,7,8]。
$b$ 的每個數字之間都差 $1$ ，而 $a$ 的前五個數字也是，所以兩個可以配對。
把 $a,b$ 都換成「比較時的數值」：

• $b[i+1]-b[i]$ 的所有值： $[1,1,1,1]$
• $a[i+1]-a[i]$ 的所有值： $[1,1,1,1,1]$

此時思路就很明顯了：用 $A$ 建立一個前綴樹，然後拿 $b$ 作為要尋找的前綴。

程式碼

1. 類別實現

1. 由於只會插入以及查詢前綴，因此變數end在此題當中毫無用處，可以直接移除。
2. 比較的數值是 $b[i+1]-b[i]$ ，不確定數值範圍，因此用哈希表去紀錄走過的路徑，而非固定大小的陣列。

#include <iostream>
#include <unordered_map>
#include <vector>
using namespace std;

struct TrieNode{
    int pass;
    // 這題不會使用到 end 變數，直接抹煞。
    unordered_map<int, TrieNode*> nexts;
};

class Trie {
private:
    TrieNode* root;
public:
    Trie() {
        root = new TrieNode();
    }
    void insert(vector<int>& nums) {
        TrieNode* cur = root;
        for(int x : nums) {
            if(cur->nexts.find(x) == cur->nexts.end()) { // 沒有創建過新的節點
                cur->nexts[x] = new TrieNode();
            }
            cur = cur->nexts[x];
            cur->pass++;
        }
    }
    int prefixNum(vector<int>& nums) {
        TrieNode* cur = root;
        for(int x : nums) {
            if(cur->nexts.find(x) == cur->nexts.end()) { // 沒有創建過的節點，無路可走
                return 0;
            }
            cur = cur->nexts[x];
        }
        return cur->pass;
    }
};

class Solution {
private:
    vector<int> getDiff(vector<int>& nums) {
        vector<int> res;
        for(int i = 1; i < nums.size(); i++) {
            res.push_back(nums[i] - nums[i - 1]);
        }
        return res;
    }
public:
    vector<int> countConsistentKeys(vector<vector<int> >& b, vector<vector<int> >& a) {
        Trie trie;
        for(auto& nums : a) {
            vector<int> diff = getDiff(nums);
            trie.insert(diff);
        }

        vector<int> res;
        for(auto& nums : b) {
            vector<int> diff = getDiff(nums);
            res.push_back(trie.prefixNum(diff));
        }

        return res;
    }
};

2. 靜態數組實現

改成靜態數組實現時，因為相減的數值範圍會很大，而且有負數，需要修改原先經典的做法。

• 假如當前的 $a = [1,19,10,3]$ ， $a[i+1]-a[i]=[18,-9,-7]$ 。
• 建立節點的方式是將數字拆成各自的位數，並且加入負號和終止符號。

用10表示負號，11表示終止符號。

※ 假如不加終止符號， $[1,8,-9,-7]$ 會被視為相同的前綴。
如果覺得特判太麻煩，直接將數字轉字串之後切分每一位也是可以的。

#include <iostream>
#include <unordered_map>
#include <vector>
#include <array>
using namespace std;

const int MX = 2e6 + 1;
int cnt;
array<array<int, 12>, MX> trie; // 0 ~ 9, 負號(10), 終止符號(11)
array<int, MX> pass;

void insert(vector<int>& nums) {
    int cur = 1;
    for(int num : nums) {
        if(num < 0) { // 處理負數
            if(trie[cur][10] == 0) {// 沒有創建過新的節點
                trie[cur][10] = ++cnt;
            }
            cur = trie[cur][10];
            pass[cur]++;
        }
        num = abs(num);
        while(num) {
            int x = num % 10;
            // 插入每一位數的節點
            if(trie[cur][x] == 0) {
                trie[cur][x] = ++cnt;
            }
            cur = trie[cur][x];
            pass[cur]++;
            num /= 10;
        }
        // 終止符號
        if(trie[cur][11] == 0) {
            trie[cur][11] = ++cnt;
        }
        cur = trie[cur][11];
        pass[cur]++;
    }
}

int prefixNum(vector<int>& nums) {
    int cur = 1;
    for(int num : nums) {
        if(num < 0) { // 處理負數
            if(trie[cur][10] == 0) {// 沒有創建過新的節點
                return 0;
            }
            cur = trie[cur][10];
        }
        num = abs(num);
        while(num) {
            int x = num % 10;
            // 插入每一位數的節點
            if(trie[cur][x] == 0) {
                return 0;
            }
            cur = trie[cur][x];
            num /= 10;
        }
        // 終止符號
        if(trie[cur][11] == 0) {
            return 0;
        }
        cur = trie[cur][11];
    }
    return pass[cur];
}

class Solution {
private:
    vector<int> getDiff(vector<int>& nums) {
        vector<int> res;
        for(int i = 1; i < nums.size(); i++) {
            res.push_back(nums[i] - nums[i - 1]);
        }
        return res;
    }
public:
    vector<int> countConsistentKeys(vector<vector<int> >& b, vector<vector<int> >& a) {
        cnt = 1;
        for(auto& nums : a) {
            vector<int> diff = getDiff(nums);
            insert(diff);
        }

        vector<int> res;
        for(auto& nums : b) {
            vector<int> diff = getDiff(nums);
            res.push_back(prefixNum(diff));
        }

        return res;
    }
};

複雜度分析

• 時間複雜度：$O(n+m)$
• 空間複雜度：$O(n)$