[Leetcode] 3583. Count Special Triplets

0보다 큰 정수 배열 nums 가 주어진다. 서로 다른 세 점의 조합 (i, j, k) 에서 0 <= i < j < k < n 이며 nums[i] === nums[j] * 2와 nums[k] === nums[j]를 만족하는 조합의 갯수를 찾는 문제.

한마디로 배열의 값이 [2a, a, 2a] 형태가 되는 조합을 찾으면 된다. 나이브하게 생각하면 O(N^3) 으로 해결할 수 있지만 배열의 길이를 10^5 로 주는 걸로 보아 O(N^2) 이상이 되면 터지도록 만들어 놓은 문제이다. 배열을 두 번 또는 세 번 돌면서 O(정수 * N) 형태로 만들어 O(N) 형태로 근사되도록 하면 된다.

먼저 배열의 어떤 원소가 어디에 있는지는 알아야하기 때문에 배열을 한번 돌면서 rigthMap에 저장한다. 그 후 인덱스를 0부터 끝까지 돌면서 해당 원소의 두배 값이 좌우에 있는지, 있다면 조합의 경우의 수를 더해주면 된다.

풀었는데 처음에 시간초과가 뜨길래 당황했는데, 시간복잡도가 아니라 콘솔로그를 지우지 않아 출력에도 시간이 걸려서 만난 시간초과였다;;; console.log 를 지워주니 바로 통과.

const MOD = 1000000007;

function specialTriplets(nums: number[]): number {
    const leftMap = new Map();
    const rightMap = new Map();

    for (const num of nums) {
        rightMap.set(num, (rightMap.get(num) || 0) + 1);
    }

    let answer = 0;

    for (let i = 0; i < nums.length; i++) {
        const j = nums[i]
        
        rightMap.set(j, rightMap.get(j) - 1)

        const jDouble = j * 2;

        const leftCount = leftMap.get(jDouble) || 0;
        const rightCount = rightMap.get(jDouble) || 0;

        if (leftCount > 0 && rightCount > 0) {
            answer = (answer + leftCount * rightCount) % MOD;
        }

        leftMap.set(j, (leftMap.get(j) || 0) + 1)
    }

    return answer;
};