[백준/8462] 배열의 힘 (Python)

Problem : https://www.acmicpc.net/problem/8462

8462번: 배열의 힘

 

Difficulty : Platinum 2

 

Status : Solved

 

Time : 00:19:13

 

---

 

문제 설명

 

 

자연수 n개로 이루어진 배열 a_1,a_2,a_3, ... ,a_n이 있다.

 

l부터 r까지 부분 배열은 a_l, ..., a_r 이다.

K_s는 부분 배열 안에 있는 자연수 s의 개수이다.

부분 배열의 힘이란 모든 자연수 s에 대해서, K_s * K_s * s를 합한 값이다.

배열과 부분 배열의 범위가 주어졌을 때, 각 부분 배열의 힘을 구하는 프로그램을 작성하시오.

 

 

입력 및 출력

 

입력

 

첫째 줄에 배열의 크기 n과 부분 배열의 개수 t가 주어진다. (1 ≤ n, t ≤ 10^5) 둘째 줄에는 n개의 자연수 a_i (1 ≤ a_i ≤ 10^6)가 주어진다. 다음 t개 줄에는 부분 배열의 범위 l_i와 r_i가 주어진다. (1 ≤ l_i ≤ r_i ≤ n)

 

출력

 

입력으로 주어지는 각 부분 배열의 힘을 출력한다.

 

입력 예시

 

8 3
4 3 1 1 1 3 1 2
2 7
1 6
3 8

 

출력 예시

 

28
25
21

 

 

---

 

풀이

 

구간 쿼리의 효율적 재활용이 필요하다. N, M의 최댓값이 매우 크기 때문에 오프라인 쿼리와 mo's 알고리즘을 써 볼 수 있겠다.

 

값의 업데이트는 더 간단하다. 추가 혹은 감소되는 수가 s, 그리고 변동 이전의 s의 개수가 k라고 두자.

 

• 수 s가 추가될 경우, k^2 * s값이 (k + 1)^2 * s로 증가한다. 즉 그 차인 (2*k+1)*s를 더해주면 된다.
• 마찬가지로 수 s가 감소할 경우 k^2 * s값이 (k-1)^2*s로 감소한다. 즉 그 차인 (2*k-1)*s만큼을 빼주면 된다.

 

풀이 코드

from collections import defaultdict
import sys
input = sys.stdin.readline

N, M = map(int, input().split())
nums = list(map(int, input().split()))
queries = [[i] + list(map(lambda x : int(x)-1, input().split())) for i in range(M)]
queries.sort(key = lambda x : (x[1] // N ** 0.5, x[2]))

num_dict = defaultdict(int)
i, j, val = 0, -1, 0
ans = [0]*M

for idx, l, r in queries :
  while j < r :
    j += 1
    val += (2*num_dict[nums[j]] + 1)*nums[j]
    num_dict[nums[j]] += 1
  while j > r :
    val -= (2*num_dict[nums[j]] - 1)*nums[j]
    num_dict[nums[j]] -= 1
    j -= 1
  while i < l :
    val -= (2*num_dict[nums[i]] - 1)*nums[i]
    num_dict[nums[i]] -= 1
    i += 1
  while i > l :
    i -= 1
    val += (2*num_dict[nums[i]] + 1)*nums[i]
    num_dict[nums[i]] += 1
  ans[idx] = val

print(*ans, sep = '\n')

풀이 완료!