[백준] 13900 - 순서쌍의 곱의 합 (Python)

문제 링크

알고리즘

• 수학

풀이

from itertools import combinations


N = int(input())
nums = list(map(int, input().split()))
combis = list(combinations(nums, 2))

res = sum(map(lambda x: x[0] * x[1], combis))

print(res)

처음엔 조합 (combination) 을 이용해서 풀었지만
정수의 개수의 범위가 (2 ≤ N ≤ 100,000) 이기 때문에 메모리 초과가 났다.

그래서 다른 방식으로 문제를 해결했다.

만약 (1, 2, 3, 4, 5) 가 주어지면 답은 다음과 같이 구할 수 있다.

이는 분배 법칙을 이용해 이렇게 바꿀 수 있다.

1(2+3+4+5) + 2(3+4+5) + 3(4+5) + 4(5)

그리고 이 규칙을 이용한 알고리즘은 다음과 같다.

1. (1+2+3+4+5) 를 더해서 변수에 넣는다.
2. 1의 변수에서 처음 숫자 1을 뺀다. 그럼 (2+3+4+5) 가 남는다.
3. 빼온 숫자 1과 (2+3+4+5)를 곱한 값을 결과 값을 담을 변수에 더한다.
4. (2+3+4+5)가 들어있는 변수에서 1 다음 값인 2를 뺀다.
5. 3과 4를 반복한다.

코드로 구현하면 아래와 같다.

코드

N = int(input())
nums = list(map(int, input().split())) # N개의 숫자를 저장

sum_nums = sum(nums) # 먼저 숫자를 다 더한다.
res = 0

for n in nums: # 이제 숫자를 하나씩 빼오면서 위의 2,3,4번을 반복한다.
    sum_nums -= n
    res += sum_nums * n

print(res)