백준 5624 - 좋은 수

문제 설명

제한사항 및 입출력

풀이

문제가 굉장히 간단명료하다. 부연설명해보자면 수열에서 Ai 이전의 3개 숫자를 합해서 Ai를 만들 수 있다면 Ai는 좋은 수이다.(예제 똑바로 안 보고 바로 앞 3개로 풀이하다가 시간낭비한 사람...) 추가적으로 같은 값을 여러 번 써도 된다(인덱스 중복 허용) -> v+v+v, v+v+z 같은 조합도 가능.

 

v = (x + y) + z 에서 이항하면 v - z = x + y의 형태로 나타낼 수 있기 때문에 앞의 두 수의 합을 set에 저장한다면 앞의 숫자들을 순회하며 O(n)에 판별할 수 있다. 코드를 확인해보자

내 코드

import sys
input = sys.stdin.readline
mis = lambda: map(int, input().split())

n = int(input())
arr = [*mis()]

seen = set()       # 지금까지 앞부분에 등장한 값들의 집합
pair_sums = set()  # 앞부분 두 수의 합으로 만들 수 있는 값들의 집합

ans = 0
for v in arr:
    # 판정: v - z = x + y가능한지?
    if any((v - j) in pair_sums for j in seen):
        ans += 1
	
    # 갱신: v를 앞부분에 편입
    # v와 앞의 모든 x를 짝지어 두 수 합을 추가
    pair_sums.update(v + x for x in seen)
    
    # v를 두 번 사용하는 경우
    pair_sums.add(v + v)
    seen.add(v)

print(ans)

앞에 나온 숫자들을 순회하며 두 수의 합들을 확인하고 현재 값을 만들 수 있는지 확인한다.

이후 앞에 나온 숫자들에 현재 숫자를 더하며 두 수의 조합을 만들어서 pair_sums에 저장해주고 현재 값 두 개를 사용하는 경우도 함께 저장해준다.

 

구현 자체는 간단해서 크게 설명할만한 부분이 없다. 로직을 생각해서 숫자 조합을 O(n)으로 판별하는 것이 중요한 문제인 것 같다. 모든 세 수의 조합을 모두 계산하기에는 O(n^3)이 나오기 때문에 n <= 5000 이라는 제한 안에서는 시간 안에 해결할 수 없다.

시간복잡도

모든 수를 순회하며(n) 해당 수의 이전 모든 값 확인(~n), 현재 수를 앞의 모든 수와 더하여 새로운 조합 만들기(~n)

-> O(n^2) 

고찰

투 포인터로도 풀이할 수 있다! 시간초과가 날 줄 알았는데 아슬아슬하게 통과한다.

import bisect
import sys
mis = lambda: map(int, input().split())
input = sys.stdin.readline

n = int(input())
arr = [*mis()]

seen = set()          # 앞부분의 값들
seen_sorted = []      # 정렬된 고유값 리스트

def has_two_sum(sorted_vals, target):
    # 같은 값을 두 번 써도 되므로 l<=r 허용
    l, r = 0, len(sorted_vals) - 1
    while l <= r:
        s = sorted_vals[l] + sorted_vals[r]
        if s == target:
            return True
        if s < target:
            l += 1
        else:
            r -= 1
    return False

ans = 0
for v in arr:
    # v = j + (x+y) 판정
    good = False
    for j in seen:
        if has_two_sum(seen_sorted, v - j):
            good = True
            break
    if good:
        ans += 1

    # v를 앞부분에 편입(정렬된 고유값 유지)
    if v not in seen:
        seen.add(v)
        bisect.insort(seen_sorted, v)

print(ans)

i번째 값 v를 처리할 때, 앞에서 나온 서로 다른 값들(집합 seen)만으로 만들어진 정렬 리스트 seen_sorted에서 어떤 j 에 대해 v - j가 두 수 합으로 만들 수 있으면 v는 좋은 수이다.

 

v - z = x + y를 투 포인터로 검사하는 로직이다. “같은 값을 여러 번 써도 된다” 조건 때문에, 두 포인터는 l <= r 로 둬서 x==y도 고려해서 풀이한다.

 

투 포인터 풀이의 시간복잡도는 다음과 같다.

U_i를 i번째 단계에서의 서로 다른 앞 값 개수, U_max를 max_i U_i (최대 고유 개수)라고 했을 때

 

 

바깥 루프 -> O(U_i) 번

매번 has_two_sum(seen_sorted, v - j) 판정 비용 -> O(U_i²) (seen 내부에서 시행되기 때문)

 

갱신하는데 bisect.insort(seen_sorted, v) → O(U_i)이므로, 최악 O(N³)의 시간복잡도가 든다.