Leetcode 2916. Subarrays Distinct Element Sum of Squares II

문제

You are given a 0-indexed integer array nums.

The distinct count of a subarray of nums is defined as:

Let nums[i..j] be a subarray of nums consisting of all the indices from i to j such that 0 <= i <= j < nums.length. Then the number of distinct values in nums[i..j] is called the distinct count of nums[i..j].
Return the sum of the squares of distinct counts of all subarrays of nums.

Since the answer may be very large, return it modulo 109 + 7.

A subarray is a contiguous non-empty sequence of elements within an array.

0-indexed 정수 배열 nums가 주어집니다.

배열 nums의 부분 배열(subarray)에 대한 서로 다른 원소의 개수(distinct count)는 다음과 같이 정의됩니다.

• nums[i..j]는 0 <= i <= j < nums.length를 만족하는 인덱스 i부터 j까지의 모든 원소로 구성된 부분 배열입니다.
• nums[i..j]에 포함된 서로 다른 값의 개수를 해당 부분 배열의 distinct count라고 합니다.

nums의 모든 부분 배열에 대해 distinct count의 제곱을 합산한 값을 반환하세요.

정답이 매우 클 수 있으므로 10^9 + 7로 나눈 나머지를 반환하세요.

참고: 부분 배열(subarray)은 원본 배열에서 연속된 하나 이상의 원소로 구성된 배열입니다.

예시

입력: nums = [1,2,1]

출력: 15

설명:

가능한 부분 배열은 총 6개입니다.

• [1]: 서로 다른 원소 1개
• [2]: 서로 다른 원소 1개
• [1]: 서로 다른 원소 1개
• [1,2]: 서로 다른 원소 2개
• [2,1]: 서로 다른 원소 2개
• [1,2,1]: 서로 다른 원소 2개

모든 부분 배열의 서로 다른 원소 개수를 제곱하여 합산하면 다음과 같습니다.

1² + 1² + 1² + 2² + 2² + 2² = 15

제한

• $1 \le nums.length \le 10^5$
• $1 \le nums[i] \le 10^5$

풀이

MOD = 10 ** 9 + 7

class SegTree:
    def __init__(self, size: int):
        self.size = size
        SEGSIZE = size * 4 + 1

        self.tree = [0] * SEGSIZE
        self.nums = [0] * SEGSIZE
        self.lazy = [0] * SEGSIZE

    def _apply(self, index: int, size: int, delta: int):
        self.tree[index] += self.nums[index] * 2 * delta + delta**2 * size
        self.nums[index] += size * delta
        self.lazy[index] += delta

    def _propagation(self, left: int, right: int, index: int):
        delta = self.lazy[index]

        if left == right or delta == 0:
            return

        mid = (left + right) // 2

        self._apply(index * 2, mid - left + 1, delta)
        self._apply(index * 2 + 1, right - mid, delta)

        self.lazy[index] = 0

    def _update(
        self, seg_left: int, seg_right: int, qur_left: int, qur_right: int, index: int
    ):
        if qur_right < seg_left or qur_left > seg_right:
            return

        if qur_left <= seg_left and seg_right <= qur_right:
            self._apply(index, seg_right - seg_left + 1, 1)
            return

        self._propagation(seg_left, seg_right, index)

        mid = (seg_left + seg_right) // 2

        self._update(seg_left, mid, qur_left, qur_right, index * 2)
        self._update(mid + 1, seg_right, qur_left, qur_right, index * 2 + 1)

        self.tree[index] = self.tree[index * 2] + self.tree[index * 2 + 1]
        self.nums[index] = self.nums[index * 2] + self.nums[index * 2 + 1]

    def update(self, left: int, right: int) -> int:
        self._update(0, self.size - 1, left, right, 1)
        return self.tree[1]


class Solution:
    def sumCounts(self, nums: List[int]) -> int:
        N = len(nums)
        seg = SegTree(N)
        seen = {}
        ans = 0

        for i, v in enumerate(nums):
            left = seen[v] + 1 if v in seen else 0
            right = i
            ans = (ans + seg.update(left, right)) % MOD
            seen[v] = i

        return ans

기초 생각

길이가 3인 배열 [0, 0, 0]을 준비하고, nums = [1, 2, 1]이 주어졌다고 생각해 보자.

이 배열의 각 인덱스는 해당 인덱스에서 시작하여 현재 인덱스에서 끝나는 부분 배열의 서로 다른 원소 개수를 나타낸다.

0번 인덱스까지 확인했을 때:

[1, 0, 0]

• [1] → 1개
• 나머지는 아직 부분 배열이 만들어지지 않았다.

1번 인덱스까지 확인했을 때:

[2, 1, 0]

• [1, 2] → 2개
• [2] → 1개

2번 인덱스까지 확인했을 때:

[2, 2, 1]

• [1, 2, 1] → 2개
• [2, 1] → 2개
• [1] → 1개

각 단계에서 배열의 모든 원소를 제곱하여 합산하면 다음과 같다.

• [1, 0, 0] → 1² = 1
• [2, 1, 0] → 2² + 1² = 5
• [2, 2, 1] → 2² + 2² + 1² = 9

따라서 정답은 1 + 5 + 9 = 15가 된다.

즉, 우리가 해야 할 일은 다음 두 가지다.

1. 새로운 원소를 확인할 때마다 필요한 구간의 값을 1 증가시킨다.
2. 갱신된 배열의 제곱합을 정답에 누적한다.

하지만 매번 배열을 직접 수정하고 전체 제곱합을 계산하면 최악의 경우 O(N²)의 시간이 소요된다.

따라서 구간 갱신과 제곱합 계산을 효율적으로 처리할 수 있는 세그먼트 트리를 사용한다.

제곱합을 빠르게 구하는 방법

다음 항등식을 살펴보자.

(x + 1)² = x² + 2x + 1

기존 값 x가 1 증가했을 때, 제곱한 값은 기존 제곱값에서 2x + 1만큼 증가한다.

이를 구간 전체에 적용할 수 있다.

세그먼트 트리의 구현

세그먼트 트리의 각 변수는 다음과 같은 역할을 한다.

1. self.tree: 해당 구간에 포함된 값들의 제곱합을 저장한다.
2. self.nums: 해당 구간에 포함된 값들의 합을 저장한다.
3. self.lazy: 자식 노드에 아직 전파하지 않은 누적 증가량을 저장한다.

여기서 tree와 nums를 동시에 관리하는 것이 핵심이다.

제곱합을 갱신하기 위해서는 기존 값들의 합이 필요하기 때문이다.

구간의 값이 증가할 때 제곱합을 빠르게 구하는 방법

어떤 구간에 n개의 원소가 있다고 가정하자.

[x₁, x₂, x₃, ..., xₙ]

이 구간의 모든 원소에 δ를 더한다면 새로운 제곱합은 다음과 같다.

(x₁ + δ)² + (x₂ + δ)² + ... + (xₙ + δ)²

각 항을 전개하면 다음과 같다.

= (x₁² + 2x₁δ + δ²) + ... + (xₙ² + 2xₙδ + δ²)

같은 종류의 항끼리 묶으면 다음과 같다.

= Σxᵢ² + 2δΣxᵢ + nδ²

우리는 세그먼트 트리에 다음 값들을 이미 저장하고 있다.

• Σxᵢ² → self.tree[index]
• Σxᵢ → self.nums[index]
• n → 구간의 길이인 size

따라서 새로운 제곱합은 다음과 같이 계산할 수 있다.

새로운 제곱합 = 기존 제곱합 + 2δ × 기존 합 + nδ²

코드로 표현하면 다음과 같다.

self.tree[index] += self.nums[index] * 2 * delta + delta**2 * size

한편, 구간의 모든 원소에 δ를 더했다면 일반 합은 nδ만큼 증가한다.

self.nums[index] += size * delta

이 두 공식을 이용하면 구간에 포함된 원소를 하나씩 방문하지 않고도 O(1)에 해당 구간의 합과 제곱합을 갱신할 수 있다.

Lazy Propagation을 사용하는 이유

앞서 구간 전체의 합과 제곱합을 빠르게 갱신하는 방법을 알아냈다.

하지만 세그먼트 트리의 어떤 구간을 갱신할 때마다 모든 자식 노드에 변경 사항을 전파한다면 여전히 많은 연산이 필요하다.

이를 해결하기 위해 Lazy Propagation을 사용한다.

핵심 아이디어는 현재 노드의 값만 먼저 갱신하고, 자식 노드에 대한 갱신은 실제로 필요할 때까지 미루는 것이다.

_apply() 함수는 다음 세 가지 작업을 수행한다.

1. tree[index]에 구간의 제곱합 변화량을 반영한다.
2. nums[index]에 구간의 일반 합 변화량을 반영한다.
3. lazy[index]에 자식 노드로 전달해야 할 증가량을 누적한다.

즉, _apply()는 현재 노드의 정보를 갱신하면서 이후 전파해야 할 변경 사항을 기록하는 역할을 한다.

반면 _propagation()은 이렇게 미뤄 둔 변경 사항을 실제 자식 노드에 전달한다.

현재 노드의 lazy[index]가 0이 아니라면, 이를 왼쪽과 오른쪽 자식에 각각 _apply()로 반영한다.

이때 각 자식이 담당하는 구간의 길이를 정확히 전달해야 한다.

전파가 완료되면 현재 노드의 lazy[index]를 0으로 초기화한다.

중요한 점은 이미 갱신하려는 구간 전체를 담당하는 노드에 도착했다면 더 이상 자식으로 내려갈 필요가 없다는 것이다.

따라서 _update()는 다음 순서로 동작한다.

1. 현재 구간과 갱신 구간이 겹치지 않으면 종료한다.
2. 현재 구간이 갱신 구간에 완전히 포함되면 _apply()만 수행하고 종료한다.
3. 일부만 겹친다면 _propagation()으로 기존 변경 사항을 자식에게 전달한다.
4. 양쪽 자식을 재귀적으로 갱신한다.
5. 자식들의 tree, nums를 합쳐 현재 노드의 값을 다시 계산한다.

이렇게 하면 구간 갱신을 O(log N)에 처리할 수 있다.

어느 구간을 갱신해야 할까?

이제 새로운 원소를 확인했을 때 어느 구간의 값을 1 증가시켜야 하는지 결정해야 한다.

현재 인덱스를 i, 현재 원소를 v라고 하자.

우리가 관리하는 배열의 인덱스 j는 nums[j..i]에 포함된 서로 다른 원소의 개수를 나타낸다.

새로운 원소 v가 등장했을 때는 두 가지 경우가 있다.

1. 이전에 등장한 적 없는 원소

현재 원소를 추가하면 이전의 모든 시작점에서 서로 다른 원소의 개수가 1 증가한다.

또한 현재 위치에서 시작하는 새로운 부분 배열도 만들어진다.

따라서 [0, i] 구간을 1 증가시키면 된다.

2. 이전에 등장했던 원소

v가 마지막으로 등장한 위치를 prev라고 하자.

시작점이 prev 이하인 부분 배열에는 이미 v가 포함되어 있으므로 서로 다른 원소의 개수가 변하지 않는다.

반대로 시작점이 prev보다 큰 부분 배열에는 v가 포함되어 있지 않았으므로 새로운 원소를 추가하면 서로 다른 원소의 개수가 1 증가한다.

따라서 [prev + 1, i] 구간만 1 증가시키면 된다.

이전 등장 위치는 seen 딕셔너리로 관리한다.

• 처음 등장한 원소 → left = 0
• 이전에 등장한 원소 → left = seen[v] + 1
• 현재 원소의 위치 → right = i

결국 seg.update(left, right)를 호출하면 필요한 부분만 갱신할 수 있다.

정답 계산

세그먼트 트리의 루트 노드인 tree[1]은 현재 인덱스에서 끝나는 모든 부분 배열의 서로 다른 원소 개수에 대한 제곱합을 저장한다.

따라서 별도의 구간 쿼리 없이 루트 노드의 값만 확인하면 된다.

각 인덱스에서 구간 갱신을 수행한 뒤 다음과 같이 정답에 누적한다.

ans = (ans + seg.update(left, right)) % MOD

모든 인덱스에 대해 이 작업을 반복하면 모든 부분 배열의 제곱합이 정확히 한 번씩 정답에 반영된다.

시간 및 공간 복잡도

시간 복잡도: O(N log N)

• 배열의 모든 원소를 한 번씩 순회한다: O(N)
• 각 원소에 대해 세그먼트 트리의 구간 갱신을 수행한다: O(log N)
• seen 딕셔너리의 조회 및 갱신은 평균 O(1)이다.

따라서 전체 시간 복잡도는 O(N log N)이다.

공간 복잡도: O(N)

• 세그먼트 트리의 tree, nums, lazy 배열에 각각 O(N)의 공간이 필요하다.
• seen 딕셔너리는 최악의 경우 O(N)의 공간을 사용한다.

따라서 전체 공간 복잡도는 O(N)이다.

핵심 정리

이 풀이의 핵심은 서로 다른 원소의 개수를 직접 계산하는 대신, 새로운 원소가 등장할 때 증가해야 하는 시작점의 범위를 찾아 갱신하는 것이다.

그리고 구간의 제곱합을 효율적으로 관리하기 위해 다음 관계식을 사용한다.

Σ(xᵢ + δ)² = Σxᵢ² + 2δΣxᵢ + nδ²

이 관계식을 이용하면 세그먼트 트리의 각 노드에 일반 합과 제곱합을 함께 저장하고, Lazy Propagation을 통해 구간 갱신을 효율적으로 처리할 수 있다.

결과적으로 모든 부분 배열을 직접 열거하는 O(N²) 풀이를 O(N log N)으로 개선할 수 있다.