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

문제 요약

• 알고리즘 분류:  구간합, 누적합, 수학
• 난이도: Medium
• 문제내용:
  • 정수 배열 arr에서 세 개의 인덱스 (i, j, k)를 선택해야 합니다.
    • 인덱스 조건: 0 ≤ i < j ≤ k < arr.length
    • XOR 조건: arr[i]^arr[i+1]^...^arr[j−1]=arr[j]^arr[j+1]^...^arr[k]이때, 아래 조건을 만족해야 합니다.
  • 이러한 (i, j, k) 조합의 개수를 구하세요.
• 사이트 주소: https://leetcode.com/problems/count-triplets-that-can-form-two-arrays-of-equal-xor/description/

문제풀이

이번 문제는 완전탐색으로 풀면 O(N^4)나오지만 그러면 시간 초과가 나서 O(N^3)으로 풀어야 되는 문제이다. 그래서 이번에는 누적합과 구간합을 응용해서 풀것이다. 구간합에 대한 이론은 아래의 사이트에 참조하면 된다.

https://jih3508.tistory.com/50

[알고리즘 이론] 구간합, 누적합(prefix sum)

문제 접근 방법

 위에 글을 보면 누적된 합으로 구간에 대한 빼기로 설명이 되어 있어서 XOR연산에 대한 구간 합에 대해서는 어떻게 접근 해야 할지 어려울수 있다. 그래서 XOR에 대한 기본 적인 개념만 알면 구간 합에 바로 응용 할 수 있다.

1. XOR 연산의 특성 활용

XOR 연산의 중요한 특성은 다음과 같다.

1. x ^ x = 0 (어떤 숫자를 자기 자신과 XOR 하면 0이 됨)
2. x ^ 0 = x (0과 XOR 하면 원래 값 그대로 유지됨)

arr이 문제에서 제공 하는 파라미터, prefix가 누적(prfix[n] = arr[1] ^ arr[2] ^ arr[3] ..... arr[n]), i < j 라고 할때

위 성질을 활용하면, arr[i] ^ arr[i+1] ...... ^ arr[j] = prfix [i] ^ prfix [j] 도 성립이 된다. 그 이유는

 prfix [i] ^ prfix [j]  = arr[1] ^ arr[2] ^ arr[3] ^ ... arr[i] ^ arr[1] ^ arr[2] ^ arr[3] ^ ... arr[i] ^ arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j] 

 위 식으로 풀 수 있다. 그러면 arr[1] ^ arr[2] ^ arr[3] ^ ... arr[i] 부분은 1번쩨 x ^ x = 0 성질로 0으로 나올 것이라서

arr[1] ^ arr[2] ^ arr[3] ^ ... arr[i] ^ arr[1] ^ arr[2] ^ arr[3] ^ ... arr[i] ^ arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j]  = 0 ^ arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j] 

 위 와 같이 나올 것이고 0 ^   arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j] 을 0 ^ x = x 상질로 아래 식이 된다.

0 ^ arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j]   = arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j]

 전체 문제 풀이는 아래와 같다.

prfix [i] ^ prfix [j]
 = arr[1] ^ arr[2] ^ arr[3] ^ ... arr[i] ^ arr[1] ^ arr[2] ^ arr[3] ^ ... arr[i] ^ arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j] 
 = (arr[1] ^ arr[2] ^ arr[3] ^ ... arr[i] ^ arr[1] ^ arr[2] ^ arr[3] ^ ... arr[i]) ^ arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j]
 = 0 ^ arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j]
 = 0 ^ (arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j])
 = arr[i + 1] ^ arr[i + 2] ^ arr[i + 3] .... ^ arr[j - 1] ^ arr[j]

 

그래서 결론은 prfix [i] ^ prfix [j]는 i + 1 부터 j까지 구간합이라고 보면 된다.

 다시 문제로 돌아가서 (i, j, k)에서 같은 개수를 구하는 것은 prfix [i] ^ prfix [j ] ,  prfix [j] ^ prfix [k + 1] 두 개가 같은 개수를 구하려는 것과 같다.

  (i, j, k) 모든 경우의 수는 i 부터 배열 크기 - 1까지 고 j는 i + 1 부터 배열 크기까지 k는 j부터 배열 크기까지 반복문 돌리면 된다.

정리 하면 아래와 같다.

1.  i 부터 배열 크기 - 1까지 고 j는 i + 1 부터 배열 크기까지 k는 j부터 배열 크기까지 반복문 돌린다.
2.  prfix [i] ^ prfix [j ] ,  prfix [j] ^ prfix [k + 1]  같은지 구한다.

 위 방식대로 구현 하면 시간 복잡도는 O(N^3)이다. 자세한 것은 아래 코드 보면 이해가 될것이다.

Code

Python

class Solution:
    def countTriplets(self, arr: List[int]) -> int:
        size = len(arr)
        prefixXor = [0 for _ in range(size + 1)]
        prefixXor[1] = arr[0]

        # prefixXor[i]는 arr[0]부터 arr[i-1]까지의 XOR 연산 결과를 저장
        for i in range(1, size):
            prefixXor[i + 1] = prefixXor[i] ^ arr[i]

        count = 0
        for i in range(size - 1):
            for j in range(i + 1, size):
                for k in range(j, size):
                    # arr[i]부터 arr[j-1]까지의 XOR 값과 arr[j]부터 arr[k]까지의 XOR 값이 같을 경우 카운트 증가
                    count += (prefixXor[j] ^ prefixXor[i]) == (prefixXor[k + 1] ^ prefixXor[j])

        return count

Java

class Solution {
    public int countTriplets(int[] arr) {

        int size = arr.length;
        int[] prefixXor = new int[size + 1];
        prefixXor[1] = arr[0];

        for (int i = 1; i < size ; i++) {
            prefixXor[i + 1] = prefixXor[i] ^ arr[i];
        }

        int count = 0;
        for (int i = 0; i < size - 1; i++) {
            for (int j = i + 1; j < size; j++) {
                for (int k = j; k < size; k++) {
                    count += (prefixXor[j] ^ prefixXor[i]) == (prefixXor[k + 1] ^ prefixXor[j])? 1 : 0;
                }
            }
        }
        
        return count;
    }
}

Javascript

var countTriplets = function(arr) {

    const size = arr.length;
    let prefixXor = Array(size).fill(0);

    prefixXor[1] = arr[0];
    // prefixXor[i]는 arr[0]부터 arr[i-1]까지의 XOR 연산 결과를 저장
    for (let i = 1; i < size; i++) {
        prefixXor[i + 1] = arr[i] ^ prefixXor[i];
    }

    let count = 0;
    for (let i = 0; i < size - 1; i++) {
        for (let j = i + 1; j < size; j++) {
            for (let k = j; k < size; k++) {
                // arr[i]부터 arr[j-1]까지의 XOR 값과 arr[j]부터 arr[k]까지의 XOR 값이 같을 경우 카운트 증가
                count += (prefixXor[j] ^ prefixXor[i]) === (prefixXor[k + 1] ^ prefixXor[j]);
            }
        }
    }

    return count;
};