백준 21818 - Do You Know Your ABCs? [완전탐색]

문제를 거꾸로 생각해보았다.

 

(A, B, C)가 있을 때 [A, B, C, A+B, A+C, B+C, A+B+C] 를 만들었다고 치면

Elsie가 말한 숫자들이 전부 위 배열 안에 있어야 한다. 

 

그렇다면 A, B, C가 될 수 있는 후보군이 뭘까?

 

우리가 찾을 수 있는 숫자들은 위에 7개밖에 없다. 

위의 숫자들을 더하거나 빼면 후보군이 나올수가 있다. 

(A+B+C) - (A+B) = C

                - (A+C) = B

                - (B+C) = A

                - (A) = B+C

                - (B) = A+C

.....

..

이렇게 숫자를 뺀다가 핵심인거 같다. 

받은 배열의 숫자의 차들이 답의 후보군이다. 

 

Elsie가 말한 배열을 L이라고 하고 (A, B, C)로 만들 수 있는 숫자들의 전체 배열을 tmp라고 해보자. 

그러면 set(tmp) - set(L)의 길이가 0이어야 한다. 

 

자바하고 c++은 귀찮게 3중 for문을 돌아야 하지만 파이썬은 combinations_with_replacement가 있다.

from itertools import combinations_with_replacement
n = int(input())
for _ in range(n):
    N = int(input())
    L, E = sorted(list(map(int, input().split()))), set()
    for i in range(N):
        E.add(L[i])
        for j in range(i+1, N):
            E.add(L[j] - L[i])
    A, L = 0, set(L)
    for a, b, c in list(combinations_with_replacement(list(E), 3)):
        tmp = set([a, b, c, a+b, a+c, b+c, a+b+c])
        if len(L - tmp) == 0: A += 1
    print(A)

set은 집합이여서 그냥 - 로도 간단하게 구현할 수 있다

파이썬은신이다

 

좀더짧게

from itertools import combinations_with_replacement, combinations
n = int(input())
for _ in range(n):
    N, L, A = int(input()), sorted(list(map(int, input().split()))), 0
    for a, b, c in list(combinations_with_replacement(list(set(i[1]-i[0] for i in combinations([0] + L, 2))), 3)):
        if len(set(L) - set([a, b, c, a+b, a+c, b+c, a+b+c])) == 0: A += 1
    print(A)