[백준 1994번] 등차수열 (Python3)

import sys,collections
input = sys.stdin.readline

N = int(input())
nums = [int(input()) for i in range(N)]
result = collections.Counter(nums).most_common(n=1)[0][1]
nums = sorted({*nums})
N = len(nums)

Index = {nums[i]:i for i in range(N)}
visited = [[0]*N for i in range(N)]

for i in range(N-result):
  for j in range(i+1,N):
    if visited[i][j]:
      continue
    diff,l = nums[j]-nums[i],2
    while Index.get(nums[j]+diff):
      k = Index[nums[j]+diff]
      visited[j][k] = 1
      l += 1
      j = k
    result = max(result,l)

print(result)

3중 반복문으로 간단하게 해결하였다.

일단 처음으로 python의 collections.Counter 라는 함수를 사용해보았다. 문제에서 공차가 0인 등차수열도 고려해주어야 한다고 나오는데, 내 풀이에 중복되는 숫자들이 방해되기도 하고, 어차피 중복되는 숫자들은 공차가 0인 등차수열에밖에 쓰이지 않기 때문에 Counter의 most_common 함수를 사용하여 가장 긴 공차가 0인 등차수열을 우선 구하고, 중복되는 숫자들을 제거한 뒤 문제를 풀었다.

 

그리고 3중 반복문으로 풀어주었는데, 당연히 그냥 쌩으로 3중 반복을 하면 시간복잡도가 O(N^3)이기 때문에 풀 수 없다. 따라서 중복을 체크해주어야하는데, 이는 visited를 사용해서 해결하였다. 중복을 제거하면 O(N^2)만에 해결할 수 있다.

 

3중 반복문의 알고리즘을 설명하면,

1. visited 배열을 2차원으로 만든다. visited[i][j]는 i번 숫자 다음에 j번 숫자가 등장하는 등차수열을 확인해보았다는 의미이다.

2. 숫자들을 정렬하고 Index dict를 만든다. 모든 숫자들의 인덱스를 저장한다.

3. 현 위치에서 방문하지 않은 모든 노드에 대해서 등차수열의 길이를 확인한다. 공차는 현 노드와 탐색하려는 노드의 차이이며 다음 숫자가 존재하는지의 여부는 dict.get 함수를 이용한다.

4. 등차수열의 길이를 확인하면서 숫자를 갱신할 때, 이전노드와 다음노드의 visited도 갱신한다.

5. 등차수열의 최대 길이를 출력한다.

 

 

오늘의 교훈) Counter 함수를 앞으로 활용해보자