leetcode之形成两个异或相等数组的三元组数目（C++）

参考链接

1. https://leetcode-cn.com/problems/count-triplets-that-can-form-two-arrays-of-equal-xor/
2. https://leetcode-cn.com/problems/count-triplets-that-can-form-two-arrays-of-equal-xor/solution/xing-cheng-liang-ge-yi-huo-xiang-deng-sh-jud0/

题目描述

给你一个整数数组 arr 。

现需要从数组中取三个下标 i、j 和 k ，其中 (0 <= i < j <= k < arr.length) 。

a 和 b 定义如下：

a = arr[i] ^ arr[i + 1] ^ … ^ arr[j - 1]
b = arr[j] ^ arr[j + 1] ^ … ^ arr[k]
注意：^ 表示 按位异或 操作。

请返回能够令 a == b 成立的三元组 (i, j, k) 的数目。

解题思路

首先前缀和求出0到数组所有位置的异或值，用$S[k+1]$表示数组下标0到$k$的异或值。所以$a=S[i]⨁S[j]$，$b=S[j]⨁S[k+1]$。要使$a=b$，只需S[i] = S[k + 1]。直接使用三重循环遍历所有情况是容易想到的，难点在于如何优化。由于$S[i]=S[k+1]$的条件与$j$无关，所以只要$S[i]=S[k+1]$，$j$可以是$i+1$到$k$间的任一值，共$k-i$种，于是优化到了二重循环。

进一步分析，若$i=i_1,i_2,\cdots i_m$都满足$S[i]=S[k+1]$，那么，$\left(k-i_1\right)+\left(k-i_2\right)+\left(k-i_m\right)=m\times k-\left(i_1+i_2+\cdots i_m\right)$。所以我们需要用两个哈希表分别存储S[k]出现的次数和值为S[k]的下标之和。甚至，我们可以将这一步与求前缀和合并，只需要一次遍历。

代码

三重循环

class Solution {
public:
    int countTriplets(vector<int>& arr) {
        vector<int> nums(arr.size() + 1);
        for (int i = 1; i <= arr.size(); i ++)
        {
            nums[i] = nums[i - 1] ^ arr[i - 1];
        }

        int res = 0;
        for(int i = 0, j = 1, k = 1; j < arr.size(); j ++)
        {
            for (i = j - 1; i >= 0; i --)
            {
                for (k = j; k < arr.size(); k ++)
                {
                    if (nums[i] == nums[k + 1])
                    {
                        res += 1;
                    }
                }
            }
        }
        return res;
    }
};

二重循环

class Solution {
public:
    int countTriplets(vector<int> &arr) {
        int n = arr.size();
        vector<int> s(n + 1);
        for (int i = 0; i < n; ++i) {
            s[i + 1] = s[i] ^ arr[i];
        }
        int ans = 0;
        for (int i = 0; i < n; ++i) {
            for (int k = i + 1; k < n; ++k) {
                if (s[i] == s[k + 1]) {
                    ans += k - i;
                }
            }
        }
        return ans;
    }
};

一重循环

class Solution {
public:
    int countTriplets(vector<int> &arr) {
        int n = arr.size();
        unordered_map<int, int> cnt, total;
        int ans = 0, s = 0;
        for (int k = 0; k < n; ++k) {
            int val = arr[k];
            if (cnt.count(s ^ val)) {
                ans += cnt[s ^ val] * k - total[s ^ val];
            }
            ++cnt[s];
            total[s] += k;
            s ^= val;
        }
        return ans;
    }
};