한 국가는 오랜 분단 상태를 관리하기 위해 국경선을 따라 여러 감시 구역을 운영하고 있다.
국경선에는 왼쪽에서 오른쪽으로 번호가 매겨진 N개의 감시 구역이 있으며, 각 구역의 긴장도는 Ai로 주어진다.
정부는 긴장도 차이가 큰 구역 쌍을 잠재적인 분단 위험 구역으로 분류한다.
두 구역 i, j (i < j)에 대해 다음 조건을 만족하면 위험 구역 쌍이라고 한다.
Aj - Ai >= K
위험도는 모든 위험 구역 쌍에 대해 구역 간 거리 (j - i)를 더한 값으로 정의된다.
국경선의 긴장도 정보가 주어질 때 전체 위험도를 계산하여라.
첫째 줄에 정수 N, K가 주어진다.
둘째 줄에 N개의 정수 A1, A2, ..., AN이 공백으로 구분되어 주어진다.
전체 위험도를 출력한다.
1 ≤ N ≤ 200,0000 ≤ Ai ≤ 10^90 ≤ K ≤ 10^9
5 4
2 7 3 10 8
6 5
1 2 8 4 12 15
13
21
위험 구역 쌍은 다음과 같다.
| 구역 쌍 | 거리 |
|---|---|
| (1, 2) | 1 |
| (1, 4) | 3 |
| (1, 5) | 4 |
| (3, 4) | 1 |
| (3, 5) | 2 |
| (2, 4) | 2 |
따라서 위험도는
1 + 3 + 4 + 1 + 2 + 2 = 13
이다.
위험 구역 쌍의 거리 합은
21
이 된다.
모든 구역 쌍을 직접 검사하면 시간복잡도는 O(N²)이므로 사용할 수 없다.
조건
Aj - Ai >= K
는
Ai <= Aj - K
와 같다.
현재 구역 j를 처리할 때 조건을 만족하는 이전 구역들의 개수를 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)