[LeetCode] 1442. Count Triplets That Can Form Two Arrays of Equal XOR

https://leetcode.com/problems/count-triplets-that-can-form-two-arrays-of-equal-xor/description/

leetcode - Count Triplets That Can Form Two Arrays of Equal XOR

문제 유형 : 부분합, 비트마스킹

문제 난이도 : Medium

 

문제

Given an array of integers arr.

We want to select three indices i, j and k where (0 <= i < j <= k < arr.length).

Let's define a and b as follows:

• a = arr[i] ^ arr[i + 1] ^ ... ^ arr[j - 1]
• b = arr[j] ^ arr[j + 1] ^ ... ^ arr[k]

Note that ^ denotes the bitwise-xor operation.

Return the number of triplets (i, j and k) Where a == b.

 

정수 배열 arr이 주어진다.

인덱스 3개를 고른다. i, j, k. 0 <= i < j <= k < arr.len을 만족해야 한다.

a와 b를 다음과 같이 규정한다:

• a = arr[i] ^ arr[i + 1] ^ ... ^ arr[j - 1]
• b = arr[j] ^ arr[j + 1] ^ ... ^ arr[k]

a==b인 삼원소쌍의 개수를 반환하시오.

 

풀이

나는 직접 모든 3원소쌍을 만들어봤는데, xor의 성질을 이용해 더 간결한 풀이가 있었다.

 

우선, 3원소쌍을 만드는 풀이는 다음과 같다:

길이 = 2부터 길이 =n까지의 범위에서 i, k쌍을 양 끝에 배치하고, 처음에 i ~ k-1까지의 xor을 a로 하다가,

사이에 j를 넣어서 j를 k-1부터 i+1까지 a와 b에 각각 xor연산을 시켜줬다.

아래 그림을 보면, 이해가 갈 것이다.

중간 과정에서, a와 b가 같은지를 계속 확인해주고, 같을 시 누적해주면 답이 나온다.

 

이제 xor의 성질을 이용한 풀이이다.

i부터 k까지의 모든 값을 전부 xor하면 0이어야 한다.

그러면, i부터 k사이의 j가 어떻든, i부터 k사이는 모두 a와 b가 같다.

 

부분 xor합을 따로 구해놓는다.

prefix[i]는 arr[0]부터 arr[i-1]까지의 xor값이다.

prefix[i]와 prefix[k]가 같다면, i부터 k까지는 모두 xor이 0이다.

즉, 3원소쌍이 k-i개 만들어진다.

 

 

코드

C++ - 3원소쌍 만들기

class Solution {
public:
    int countTriplets(vector<int>& arr) {
        int n = arr.size();
        int ans = 0;
        int length = 2;

        while(length <= n) {
            for(int start = 0; start <= n - length; start++) {
                int a = 0, b, i;
                for(i = start; i < start + length - 1; i++) {
                    a ^= arr[i];
                }
                b = arr[i];
                if(a == b) ans++;
                for(i = i-1; i > start; i--) {
                    a ^= arr[i];
                    b ^= arr[i];
                    if(a == b) ans++;
                }
            }
            length++;
        }

        return ans;
    }
};

 

C++ - xor의 성질을 이용한 풀이

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

        for(int i = 0; i < n; i++) {
            for(int k = i+1; k < n; k++) {
                if(prefix[i] == prefix[k+1]) ans += k-i;
            }
        }

        return ans;
    }
};