[백준 시리즈] 13900번 : 실버4 - 순서쌍의 곱의 합 / Python

제목 : 순서쌍의 곱의 합

출처 : 백준 /  13900번: 순서쌍의 곱의 합 (acmicpc.net)

사용 자료구조&알고리즘 : 수학

 

문제 설명

N개의 정수 중 서로 다른 위치의 두 수를 뽑는 모든 경우의 두 수의 곱의 합을 구하라.

 

제한

시간 제한 : 1초

메모리 제한 : 256 MB

N (2<=N<=100,000)

입력 받는 각 수는 0이상 10,000 이하의 정수이다.

 

문제 풀이

1단계 : 최대 시간복잡도는 얼마인가?

N의 크기가 100000 이하이기 때문에 O(NlogN) 이하의 시간복잡도가 요구된다.

 

2단계 : naive하게 생각해보자.

이중 반복문을 이용해 모든 두 수를 뽑고 곱하고 더하면 시간복잡도가 $O(N^2)$이다.

시간 초과가 발생하며 다른 방법을 생각해내야한다.

 

3단계 : 핵심 아이디어는 무엇인가?

완전 제곱식을 이용한다.

세 개의 항에 대해 완전 제곱식은 $(a+b+c)^2 = a^2 + b^2 + c^2 + 2ab + 2bc + 2ca$ 이며, 이를 일반화 하면 아래 공식처럼 된다.

각 항의 합과 각 항의 제곱의 합만이 필요하므로 O(N) 만에 문제를 해결할 수 있다.

 

정답 코드

n = int(input())
arr = list(map(int, input().split()))

a_sum, a_squared_sum = 0, 0
for i in range(n):
    a_sum += arr[i]
    a_squared_sum += arr[i]**2

print((a_sum**2 - a_squared_sum)//2)