한 공연장에는 일렬로 배치된 N개의 마이크가 있다.
각 마이크의 민감도는 Ai로 주어진다.
두 마이크 i, j (i < j)에 대해 다음 조건을 만족하면 하울링이 발생할 수 있다고 정의한다.
Aj - Ai >= K
회사는 하울링 위험도를 측정하기 위해 모든 하울링 발생 가능 쌍에 대해 마이크 간 거리의 합을 계산하려고 한다.
하울링 위험도는 다음과 같이 정의된다.
Σ(j - i)
단, Aj - Ai >= K를 만족하는 모든 쌍 (i, j)에 대해서만 계산한다.
프로젝트에 주어진 마이크 민감도 정보가 주어질 때 하울링 위험도를 구하여라.
첫째 줄에 정수 N, K가 주어진다.
둘째 줄에 N개의 정수 A1, A2, ..., AN이 공백으로 구분되어 주어진다.
하울링 위험도를 출력한다.
1 ≤ N ≤ 200,0000 ≤ Ai ≤ 10^90 ≤ K ≤ 10^9
5 3
1 5 2 7 6
13
조건을 만족하는 마이크 쌍은 다음과 같다.
| 마이크 쌍 | 거리 |
|---|---|
| (1, 2) | 1 |
| (1, 4) | 3 |
| (1, 5) | 4 |
| (2, 4) | 2 |
| (3, 4) | 1 |
| (3, 5) | 2 |
따라서 하울링 위험도는
1 + 3 + 4 + 2 + 1 + 2 = 13
이다.
브루트포스로 모든 쌍을 확인하면 O(N²)이 되어 최대 입력을 처리할 수 없다.
조건
Aj - Ai >= K
는
Ai <= Aj - K
와 동일하다.
현재 원소 Aj를 처리할 때 조건을 만족하는 이전 원소들의 개수를 cnt, 인덱스 합을 sumIdx라고 하면 기여도는
Σ(j - i)
= cnt × j - sumIdx
가 된다.
값 기준으로 정렬한 뒤 투 포인터를 이용하여 현재 값보다 K 이상 작은 원소들을 관리하면 모든 기여도를 효율적으로 계산할 수 있다.
정렬에 O(N log N)이 필요하며 포인터는 전체 배열을 한 번만 이동하므로 전체 시간복잡도는 O(N log N)이다.
import sys
input = sys.stdin.readline
N, K = map(int, input().split())
A = list(map(int, input().split()))
arr = sorted((A[i], i + 1) for i in range(N))
answer = 0
count = 0
sum_idx = 0
left = 0
for value, idx in arr:
while left < N and arr[left][0] <= value - K:
count += 1
sum_idx += arr[left][1]
left += 1
answer += count * idx - sum_idx
print(answer)