[Programmers] Lv2. 우박수열 정적분 Python 문제 풀이

문제

https://school.programmers.co.kr/learn/courses/30/lessons/134239

프로그래머스

 

어떻게 풀었는가

문제를 세 단계로 나눠서 접근했다.

1단계: 우박수열 생성

콜라츠 추측 규칙대로 초항 k에서 시작해 1이 될 때까지 수열을 만든다. 짝수면 2로 나누고, 홀수면 3을 곱하고 1을 더하면 된다. 예를 들어 k = 5라면 [5, 16, 8, 4, 2, 1]이 된다.

2단계: 누적합으로 면적 전처리

꺾은선 그래프 아래 면적을 구하는 건 결국 인접한 두 점 사이의 사다리꼴 넓이를 구하는 것과 같다. x 구간의 너비가 항상 1이므로 x = i ~ i+1 구간의 넓이는 (seq[i] + seq[i+1]) / 2로 간단하게 계산된다.

여기서 핵심은 매 쿼리마다 구간 넓이를 처음부터 더하면 비효율적이니, 누적합(prefix sum) 배열을 미리 만들어두는 것이다. prefix[i]에 x = 0부터 x = i까지의 면적을 저장해두면, 임의의 구간 [a, end]의 면적은 prefix[end] - prefix[a]로 O(1)에 구할 수 있다.

3단계: 구간 변환 및 예외 처리

문제에서 구간 [a, b]의 b는 0 이하의 수로, 실제 끝점은 n + b이다 (n은 수열이 1에 도달하기까지의 횟수). 예를 들어 n = 5이고 b = -2라면 실제 끝점은 3이 된다. 시작점이 끝점보다 큰 경우는 유효하지 않으므로 -1을 반환한다.

시간복잡도

우박수열의 길이를 L, 쿼리 수를 Q라 하면 수열 생성과 누적합 전처리에 O(L), 각 쿼리 응답에 O(1)이므로 전체 O(L + Q)이다.

 

코드

def solution(k, ranges):
    # 1. 우박수열 생성
    seq = [k]
    while k != 1:
        if k % 2 == 0:
            k //= 2
        else:
            k = k * 3 + 1
        seq.append(k)

    n = len(seq) - 1  # 수열 길이 - 1 = 마지막 x좌표

    # 2. 각 구간 사다리꼴 넓이의 누적합
    prefix = [0] * (n + 1)
    for i in range(n):
        prefix[i + 1] = prefix[i] + (seq[i] + seq[i + 1]) / 2

    # 3. 각 range에 대해 정적분
    answer = []
    for a, b in ranges:
        end = n + b          # b는 0 이하이므로 실제 끝점 = n + b
        if a > end:          # 시작점이 끝점보다 크면 유효하지 않음
            answer.append(-1.0)
        else:
            answer.append(prefix[end] - prefix[a])

    return answer