Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 부분 배열이 등차수열로 재배열 가능한지 확인하는 프로그램

문제 개요

숫자로 이루어진 배열 nums와 두 개의 배열 l, r이 있다고 가정해 봅시다. 이때 lr은 각각 [l[i], r[i]] 형태의 범위 쿼리를 나타냅니다. 우리가 구해야 할 것은 불리언 배열 ans로, 부분 배열 nums[l[i]], nums[l[i]+1], ..., nums[r[i]]를 재배열하여 등차수열을 만들 수 있다면 ans[i]는 True, 그렇지 않다면 False가 됩니다.

여기서 등차수열이란 최소 두 개 이상의 요소로 구성되며, 인접한 두 요소 사이의 차이(공차)가 모두 동일한 수열을 의미합니다. 예를 들어 [2, 4, 6, 8, 10], [5, 5, 5, 5], [4, -2, -8, -14]는 등차수열에 해당하지만, [2, 2, 3, 6, 9]는 요소 간 차이가 일정하지 않으므로 등차수열이 아닙니다.

예시

입력이 다음과 같다고 가정해 보겠습니다.

nums = [6,8,7,11,5,9], l = [0,0,2], r = [2,3,5]

이 경우 출력은 [True, False, True]가 됩니다. 그 이유는 다음과 같습니다.

  • 쿼리 [0, 2]: 해당 수열은 [6, 8, 7]이며, [6, 7, 8]로 재배열하면 공차가 1인 등차수열이 되므로 유효합니다.

  • 쿼리 [0, 3]: 해당 수열은 [6, 8, 7, 11]이며, 어떻게 재배열하더라도 등차수열을 만들 수 없습니다.

  • 쿼리 [2, 5]: 해당 수열은 [7, 11, 5, 9]이며, [5, 7, 9, 11]로 재배열하면 공차가 2인 등차수열이 되므로 유효합니다.

해결 접근 방법

핵심 아이디어는 간단합니다. 배열을 오름차순으로 정렬한 뒤, 인접한 요소들 간의 차이가 모두 같은지 확인하면 됩니다. 만약 차이 값들이 하나로 유일하다면 해당 부분 배열은 등차수열로 재배열 가능한 것입니다. 문제 해결 과정은 다음과 같습니다.

  • l의 길이와 같은 크기의 불리언 리스트 new를 생성하고 모든 값을 True로 초기화합니다.

  • i를 0부터 l의 길이 - 1까지 반복합니다.

    • data := nums에서 인덱스 l[i]부터 r[i]까지의 부분 배열을 추출합니다.

    • data를 오름차순으로 정렬합니다.

    • d := 새로운 빈 리스트를 생성합니다.

    • j를 0부터 data의 길이 - 2까지 반복하면서 data[j+1] - data[j] 값을 d의 끝에 추가합니다.

    • d := set(d)를 리스트로 변환하여 중복된 차이 값을 제거합니다.

    • d의 크기가 1이 아니라면, 즉 서로 다른 차이 값이 존재한다면 new[i]를 False로 설정합니다.

  • 모든 반복이 끝나면 new를 반환합니다.

구현 예제

아래의 Python 구현 예제를 통해 더 자세히 이해해 보겠습니다.

def solve(nums, l, r):
    new = [True]*len(l)

    for i in range(len(l)):
        data = nums[l[i]:r[i]+1]
        data.sort()

        d = []
        for j in range(len(data) - 1):
            d.append(data[j+1] - data[j])

        d = list(set(d))
        if len(d) != 1:
            new[i] = False
    return new

nums = [6,8,7,11,5,9]
l = [0,0,2]
r = [2,3,5]
print(solve(nums, l, r))

입력

[6,8,7,11,5,9], [0,0,2], [2,3,5]

출력

[True,False,True]

복잡도 분석

각 쿼리마다 부분 배열을 정렬해야 하므로, 쿼리의 개수를 m, 배열의 최대 길이를 n이라 할 때 전체 시간 복잡도는 O(m × n log n)입니다. 공간 복잡도는 각 쿼리마다 임시 리스트를 생성하므로 O(n)입니다. 쿼리 개수가 매우 많거나 배열이 큰 경우에는 세그먼트 트리 등의 고급 자료구조를 활용한 최적화를 고려할 수 있습니다.