Programming/백준

백준 20922 - 겹치는 건 싫어 (파이썬)

pental 2025. 2. 27. 00:07

https://www.acmicpc.net/problem/20922

풀이

진짜 단순히 배열에서 K개 이상의 dic을 만들어서 종료하는 구문으로 풀어볼까 했는데, 어림도 없지 당연히 틀렸다.

import sys
input = sys.stdin.readline

N, K = map(int, input().split())
arr = list(map(int, input().split()))

dic = {}
answer = 0
for i in arr :
    if i in dic :
        dic[i] += 1
    else :
        dic[i] = 1
    
    if dic[i] > K :
        break
    else :
        answer += 1

문제를 이해하면 다음과 같다. 주어진 정수 배열 A에서 어떤 숫자도 최대 K번 이하로 등장하는 가장 긴 연속 부분 수열의 길이를 구하는 문제이다.

같은 숫자가 K + 1 번 등장하면 안된다.

이를 해결하기 위해서 투포인터 기법을 사용한다.

  • 투 포인터 알고리즘을 활용해서 start와 end를 움직이며 최적 부분 수열을 찾는다.
  • 해시 리스트를 이용해서 각 숫자의 등장 횟수를 기록한다.
  • good 변수를 활용해서 현재 상태가 유효한지를 판단한다.
  • end를 확장하면서 조건을 유지하다가, 조건이 깨지면 start를 이동하여 다시 조건을 만족하도록 조정한다.
count = [0] * 100001

숫자의 등장 횟수를 저장할 배열

count[A[start]] += 1
max_length = 0
good = True

첫번째 숫자는 1이기 때문에 count[A[start]] += 1을 통해서 첫번째 숫자를 포함하도록 진행한다.

그리고 현재 상태가 유효한지 판단하는 변수 good을 선언한다.

if good:
  max_length = max(max_length, end - start + 1)  # 최대 길이 갱신
  
  if end == N - 1:  # `end`가 배열 끝에 도달하면 종료
      break
  
  end += 1  # `end` 증가
  count[A[end]] += 1  # 새로운 숫자 추가

  if count[A[end]] == K + 1:  # 숫자가 K+1번 등장하면 조건 위반
      good = False  # `good`을 False로 설정하여 `start`를 조정하도록 만듦

반복을 진행하면서 max_length에 최대 길이를 갱신하고, end가 배열 끝에 도달하면 종료한다.

배열 끝에 도달하지 않으면 end를 1 증가하고, 새로운 숫자를 count, 즉 해시리스트에 추가한다.

그리고 숫자가 K + 1번 등장하면 조건을 위반하기 때문에 good을 False로 설정해서 start를 움직이도록 조정하는 것이다.

else:
  start += 1  # `start` 이동하여 범위를 줄임
  count[A[start - 1]] -= 1  # 제외한 숫자의 개수 감소

  if count[A[start - 1]] == K:  # 다시 유효한 범위가 되면 `good`을 True로 설정
      good = True

즉 위 코드와 같이, good이 False일 때 start를 이동해서 범위를 줄인다.

뿐만아니라, 제외한 숫자의 개수를 감소시키고, 다시 유효한 범위가 되면 good을 True로 설정한다.

시간복잡도 분석

start와 end는 한 방향으로만 움직이므로, O(N)번만 이동한다.

따라서 전체 알고리즘의 시간 복잡도는 O(N)을 가진다.

코드

# 백준 20922 - 겹치는 건 싫어
# 분류 : 투 포인터

N, K = map(int, input().split())
A = list(map(int, input().split()))

count = [0] * 100001

start = 0
end = 0

count[A[start]] += 1

max_length = 0
good = True
while True :
    if good : 
        max_length = max(max_length, end - start + 1)
        
        if end == N - 1 :
            break

        end += 1
        count[A[end]] += 1
        if count[A[end]] == K + 1 :
            good = False

    else :
        start += 1
        count[A[start - 1]] -= 1

        if count[A[start -1]] == K :
            good = True

print(max_length)