[코딩테스트] 등차수열을 이루는 연속된 부분 배열의 길이 찾기

문제

문제: 주어진 정수 배열 arr에서, 등차수열을 이루는 가장 긴 연속된 부분 배열의 길이를 찾으세요.
(2 <= len(arr) <= 있었는데 기억안남)

예시1:
arr = [2, 3, 4, -1, -6, -11, 1]
정답: 4 (4, -1, -6, -11)

예시2:
arr = [0, 0, 0, 0, 0]
정답: 5 (0, 0, 0, 0, 0)

 

내 풀이

def solution(arr):
  answer = 0
  count = 0
  temp = 0
  for i in range(1, len(arr)):
    m = arr[i] - arr[i-1]
    if count == 0 or temp == m:
      count += 1
      temp = m
    elif count != m:
      if answer < count:
        answer = count
      count = 1
      temp = m
    
  if answer < count:
    answer = count

  return answer + 1

등차수열은 연속하는 두 항의 차이가 모두 일정한 수열을 뜻하니

1번 인덱스부터 끝 인덱스까지 이전 인덱스와 비교하며 차이가 일정한지를 체크 했다.

 

차이를 임시 변수로 저장해놓고 다음 차이를 비교하며 차이가 일정한지 체크하면서

같다면 count를 올리고 만약 다르다면 다시 count를 처음부터 세는 방법을 사용했는데

다시 count를 세기 시작할 때 지금까지 셌던 count가 최대인지 확인하는 방법을 추가했다.

 

for문이 끝나고 제일 마지막까지 일정한 수열일 경우를 위해

지금까지 카운팅한 count를 다시 최대값인지 확인하는 과정을 추가

 

그 후 마지막으로 배열의 길이를 리턴해야 하니 + 1을 한 후 리턴

 

시간복잡도

정확한 시간복잡도는 O(15N-9)가 나오지만 상수를 제외한다면 O(N)이 된다.

 

풀이 개선

def solution(arr):
  answer = 2
  count = 2
  temp = arr[1] - arr[0]
  for i in range(2, len(arr)):
    if (arr[i] - arr[i-1]) == temp:
      count += 1
    else:
      temp = arr[i] - arr[i-1]
      count = 2
    answer = max(answer, count)

  return answer

시간복잡도면에선 차이가 없지만

조건을 활용해 코드를 좀 더 가독성 좋게 수정할 수 있었다.

 

우선, arr의 길이는 2 이상인 것을 생각해 애초부터 temp를 0번째, 1번째의 차이로 구해 시작할 수 있었고

덕분에 for문도 2번째 인덱스부터 시작하며 비교가 가능해졌다.

또한, temp를 갖고 시작하니 비교할 때도 m을 따로 만들 필요가 없어진데다 if문도 깔끔해짐

 

최대값을 정할 때 if문을 썼는데 파이썬의 max() 함수를 통해 간단하게 최대값을 얻을 수 있었다. 

https://wiki.python.org/moin/TimeComplexity

위 파이썬 공식문서에 따르면 max의 시간복잡도는 O(N)이라고 하니 두 수를 비교할 땐 O(2)이 될 것이다.

 

이전에는 for문이 끝나면 마지막까지 일정한 등차수열일 경우를 위해 count와 최대값 answer를 비교해주는 로직이 필요했는데

for문 안에서 count에 이미 집계되고 있기 때문에 다음으로 넘어가기 전에 최대값을 체크하는 방법을 통해

마찬가지로 생략이 가능해졌다.

 

후기

코딩테스트 경험이 별로 없기에 막상 문제를 처음 접했을 때 생각보다 풀이 방법을 이끌어내는데 시간이 오래 걸렸었는데

별다른 자료구조의 사용 없이 구현하는 문제라 더 그랬던 것 같다.

 

만약 내가 푼 방법 말고 다른 효율 좋은 방법이 있다면 그게 무엇인지 알고 싶다.

 

이와 비슷한 문제로

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

11053번: 가장 긴 증가하는 부분 수열

이런걸 찾았는데 시간 있을 때 꼭 도전해볼 것