【每日算法】詳解幾種「枚舉 & 計數」的基本思路 ｜Python 主題月

本文正在參加「Python主題月」，詳情查看 活動連接

題目描述

這是 LeetCode 上的 1711. 大餐計數 ，難度爲 中等。

Tag : 「哈希表」、「位運算」

大餐 是指 剛好包含兩道不一樣餐品 的一餐，其美味程度之和等於 2 的冪。

你能夠搭配 任意 兩道餐品作一頓大餐。

給你一個整數數組 deliciousness ，其中 deliciousness[i] 是第 i 道餐品的美味程度，返回你能夠用數組中的餐品作出的不一樣 大餐 的數量。結果須要對 $10^9$ + 7 取餘。

注意，只要餐品下標不一樣，就能夠認爲是不一樣的餐品，即使它們的美味程度相同。 

示例 1：

輸入：deliciousness = [1,3,5,7,9]

輸出：4

解釋：大餐的美味程度組合爲 (1,3) 、(1,7) 、(3,5) 和 (7,9) 。
它們各自的美味程度之和分別爲 4 、8 、8 和 16 ，都是 2 的冪。
複製代碼

示例 2：

輸入：deliciousness = [1,1,1,3,3,3,7]

輸出：15

解釋：大餐的美味程度組合爲 3 種 (1,1) ，9 種 (1,3) ，和 3 種 (1,7) 。
複製代碼

提示：

• 1 <= deliciousness.length <= $10^5$
• 0 <= deliciousness[i] <= $2^{20}$

---

枚舉前一個數（TLE）

一個樸素的想法是，從前日後遍歷 $deliciousness$ 中的全部數，當遍歷到下標 $i$ 的時候，回頭檢查下標小於 $i$ 的數是否可以與 $deliciousness[i]$ 相加造成 $2$ 的冪。

這樣的作法是 $O(n^2)$ 的，防止一樣的數值被重複計算，咱們可使用「哈希表」記錄某個數出現了多少次，但這並不改變算法仍然是 $O(n^2)$ 的。

並且咱們須要一個 check 方法來判斷某個數是否爲 $2$ 的冪：

• 樸素的作法是對 $x$ 應用試除法，固然由於精度問題，咱們須要使用乘法實現試除；
• 另外一個比較優秀的作法是利用位運算找到符合「大於等於 $x$」的最近的 $2$ 的冪，而後判斷是否與 $x$ 相同。

兩種作法差距有多大呢？方法一的複雜度爲 $O(\log{n})$，方法二爲 $O(1)$。

根據數據範圍 $0 <= deliciousness[i] <= 2^{20}$，方法一最多也就是執行不超過 $22$ 次循環。

顯然，採用何種判斷 $2$ 的冪的作法不是關鍵，在 OJ 斷定上也只是分別卡在 $60/70$ 和 $62/70$ 的 TLE 上。

但經過這樣的分析，咱們能夠發現「枚舉前一個數」的作法是與 $n$ 相關的，而枚舉「可能出現的 $2$ 的冪」則是有明確的範圍，這引導出咱們的解法二。

Java 代碼：

class Solution {
    int mod = (int)1e9+7;
    public int countPairs(int[] ds) {
        int n = ds.length;
        long ans = 0;
        Map<Integer, Integer> map = new HashMap<>();
        for (int i = 0; i < n; i++) {
            int x = ds[i];
            for (int other : map.keySet()) {
                if (check(other + x)) ans += map.get(other);
            }
            map.put(x, map.getOrDefault(x, 0) + 1);
        }
        return (int)(ans % mod);
    }
    boolean check(long x) {
        // 方法一
        // long cur = 1;
        // while (cur < x) {
        // cur = cur * 2;
        // }
        // return cur == x;
        
        // 方法二
        return getVal(x) == x;
    }
    long getVal(long x) {
        long n = x - 1;
        n |= n >>> 1;
        n |= n >>> 2;
        n |= n >>> 4;
        n |= n >>> 8;
        n |= n >>> 16;
        return n < 0 ? 1 : n + 1;
    }
}
複製代碼

Python3 代碼：

class Solution:
    mod = 10 ** 9 + 7
    def countPairs(self, deliciousness: List[int]) -> int:
        n = len(deliciousness)
        ans = 0
        hashmap = Counter()
        for i in range(n):
            x = deliciousness[i]
            for other in hashmap:
                if self.check(other+x):
                    ans += hashmap[other]
            hashmap[x] += 1
        return ans % self.mod
    
    def check(self, x):
        """ # 方法一 cur = 1 while cur < x: cur *= 2 return cur == x """
        
        # 方法二
        return self.getVal(x) == x
    
    def getVal(self, x):
        n = x - 1
        # java中 >>>：無符號右移。不管是正數仍是負數，高位統統補0。 Python不須要
        n |= n >> 1
        n |= n >> 2
        n |= n >> 4
        n |= n >> 8
        n |= n >> 16
        return 1 if n < 0 else n + 1
複製代碼

• 時間複雜度： $O(n^2)$
• 空間複雜度： $O(n)$

---

枚舉 2 的冪（容斥原理）

根據對樸素解法的分析，咱們能夠先使用「哈希表」對全部在 $deliciousness$ 出現過的數進行統計。

而後對於每一個數 $x$，檢查全部可能出現的 $2$ 的冪 $i$，再從「哈希表」中反查 $t = i - x$ 是否存在，並實現計數。

一些細節：若是哈希表中存在 $t = i - x$，而且 $t = x$，這時候方案數應該是 $(cnts[x] - 1) * cnts[x]$；其他通常狀況則是 $cnts[t] * cnts[x]$。

同時，這樣的計數方式，咱們對於二元組 $(x, t)$ 會分別計數兩次（遍歷 $x$ 和 遍歷 $t$），所以最後要利用容斥原理，對重複計數的進行減半操做。

Java 代碼：

class Solution {
    int mod = (int)1e9+7;
    int max = 1 << 22;
    public int countPairs(int[] ds) {
        Map<Integer, Integer> map = new HashMap<>();
        for (int d : ds) map.put(d, map.getOrDefault(d, 0) + 1);
        long ans = 0;
        for (int x : map.keySet()) {
            for (int i = 1; i < max; i <<= 1) {
                int t = i - x;
                if (map.containsKey(t)) {
                    if (t == x) ans += (map.get(x) - 1) * 1L * map.get(x);
                    else ans += map.get(x) * 1L * map.get(t);
                }
            }
        }
        ans >>= 1;
        return (int)(ans % mod);
    }
}
複製代碼

Python3 代碼：

class Solution:
    mod = 10 ** 9 + 7
    maximum = 1 << 22
    def countPairs(self, deliciousness: List[int]) -> int:
        hashmap = Counter(deliciousness)
        ans = 0
        for x in hashmap:
            i = 1
            while i < self.maximum:
                t = i - x
                if t in hashmap:
                    if t == x:
                        ans += (hashmap[x] - 1) * hashmap[x]
                    else:
                        ans += hashmap[x] * hashmap[t]
                i <<= 1
        ans >>= 1
        return ans % self.mod
複製代碼

• 時間複雜度：根據數據範圍，令 $C$ 爲 $2^{21}$。複雜度爲 $O(n * \log{C})$
• 空間複雜度： $O(n)$

---

枚舉 2 的冪（邊遍歷邊統計）

固然，咱們也能夠採起「一邊遍歷一邊統計」的方式，這樣取餘操做就能夠放在遍歷邏輯中去作，也就順便實現了不使用 $long$ 來計數（以及不使用 % 實現取餘）。

Java 代碼：

class Solution {
    int mod = (int)1e9+7;
    int max = 1 << 22;
    public int countPairs(int[] ds) {
        Map<Integer, Integer> map = new HashMap<>();
        int ans = 0;
        for (int x : ds) {
            for (int i = 1; i < max; i <<= 1) {
                int t = i - x;
                if (map.containsKey(t)) {
                    ans += map.get(t);
                    if (ans >= mod) ans -= mod;
                }
            }
            map.put(x, map.getOrDefault(x, 0) + 1);
        }
        return ans;
    }
}
複製代碼

Python3 代碼：

class Solution:
    mod = 10 ** 9 + 7
    maximum = 1 << 22
    def countPairs(self, deliciousness: List[int]) -> int:
        hashmap = defaultdict(int)
        ans = 0
        for x in deliciousness:
            i = 1
            while i < self.maximum:
                t = i - x
                if t in hashmap:
                    ans += hashmap[t]
                    if ans >= self.mod:
                        ans -= self.mod
                i <<= 1
            hashmap[x] += 1
        return ans
複製代碼

• 時間複雜度：根據數據範圍，令 $C$ 爲 $2^{21}$。複雜度爲 $O(n * \log{C})$
• 空間複雜度： $O(n)$

---

最後

這是咱們「刷穿 LeetCode」系列文章的第 No.1711 篇，系列開始於 2021/01/01，截止於起始日 LeetCode 上共有 1916 道題目，部分是有鎖題，咱們將先把全部不帶鎖的題目刷完。

在這個系列文章裏面，除了講解解題思路之外，還會盡量給出最爲簡潔的代碼。若是涉及通解還會相應的代碼模板。

爲了方便各位同窗可以電腦上進行調試和提交代碼，我創建了相關的倉庫：github.com/SharingSour… 。

在倉庫地址裏，你能夠看到系列文章的題解連接、系列文章的相應代碼、LeetCode 原題連接和其餘優選題解。