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

Python으로 길이가 정확히 M인 1 그룹이 존재하는 마지막 단계 찾기

문제 소개

1부터 n까지의 숫자로 이루어진 순열 배열 arr와 크기 n의 이진 문자열이 주어집니다. 이진 문자열의 모든 비트는 처음에 0으로 설정되어 있습니다. 각 단계 i(1부터 n까지, 인덱스는 1부터 시작)마다 arr[i] 위치의 비트를 1로 설정합니다. 여기에 값 m이 추가로 주어지며, 우리가 찾아야 하는 것은 길이가 정확히 m인 1 그룹이 존재하는 가장 늦은 단계입니다.

여기서 '1 그룹'이란 양쪽 어느 방향으로도 더 확장할 수 없는, 연속된 1로 이루어진 부분 문자열을 의미합니다. 만약 조건을 만족하는 그룹을 찾을 수 없다면 -1을 반환해야 합니다.

예시로 이해하기

arr = [3,5,1,2,4], m = 3이 주어진 경우를 살펴보겠습니다. 이때 출력은 4가 됩니다. 초기 이진 문자열은 "00000"이며, 단계별 변화 과정은 다음과 같습니다.

  • 1단계: "00100" → 그룹: ["1"]
  • 2단계: "00101" → 그룹: ["1", "1"]
  • 3단계: "10101" → 그룹: ["1", "1", "1"]
  • 4단계: "11101" → 그룹: ["111", "1"]
  • 5단계: "11111" → 그룹: ["11111"]

4단계에서 길이가 정확히 3인 그룹 "111"이 나타나지만, 5단계에서는 모든 비트가 하나의 그룹으로 합쳐져 길이가 5가 됩니다. 따라서 정답은 4입니다.

해결 전략

매 단계마다 문자열 전체를 다시 검사하는 대신, 두 개의 보조 배열을 활용하면 O(n) 시간 안에 문제를 해결할 수 있습니다.

  • l 배열: 각 위치를 기준으로 왼쪽에 연속해서 이어진 1의 개수를 저장
  • r 배열: 각 위치를 기준으로 오른쪽에 연속해서 이어진 1의 개수를 저장
  • num 변수: 현재 길이가 정확히 m인 그룹의 개수를 실시간으로 추적

새로운 위치에 비트를 1로 설정하면 왼쪽 그룹과 오른쪽 그룹이 하나로 합쳐집니다. 이때 다음 규칙을 적용합니다.

  • 합치기 전에 좌우 그룹 중 크기가 m이었던 그룹이 있다면 num을 1씩 감소시킵니다. 해당 그룹이 소멸되기 때문입니다.
  • 새로 만들어진 그룹의 크기(1 + 왼쪽 길이 + 오른쪽 길이)가 m이라면 num을 1 증가시킵니다.
  • num이 0보다 크면, 즉 길이가 m인 그룹이 존재하면 현재 단계를 정답 후보로 갱신합니다.

알고리즘 단계

  1. n := arr의 크기로 설정
  2. num := 0, ans := -1로 초기화
  3. l := 크기 n의 배열(0으로 채움), r := 크기 n의 배열(0으로 채움)
  4. i를 0부터 n-1까지 반복하며 다음을 수행:
    • cur := 1
    • idx := arr[i] - 1
    • r[idx]가 m과 같으면 num -= 1
    • l[idx]가 m과 같으면 num -= 1
    • cur := cur + l[idx] + r[idx]
    • cur가 m과 같으면 num += 1
    • num > 0이면 ans := max(ans, i + 1)
    • idx - l[idx] > 0이면 r[idx - l[idx] - 1] := cur
    • idx + r[idx] < n - 1이면 l[idx + r[idx] + 1] := cur
  5. ans 반환

Python 구현 코드

def solve(arr, m):
    n = len(arr)
    num = 0
    ans = -1
    l = [0] * n
    r = [0] * n
    for i in range(n):
        cur = 1
        idx = arr[i] - 1
        if r[idx] == m:
            num -= 1
        if l[idx] == m:
            num -= 1
        cur += l[idx] + r[idx]
        num += cur == m
        if num > 0:
            ans = max(ans, i + 1)
        if idx - l[idx] > 0:
            r[idx - l[idx] - 1] = cur
        if idx + r[idx] < n - 1:
            l[idx + r[idx] + 1] = cur
    return ans

arr = [3,5,1,2,4]
m = 3
print(solve(arr, m))

입력

[3,5,1,2,4], 3

출력

4

복잡도 분석

이 알고리즘은 각 위치를 한 번씩만 처리하므로 시간 복잡도는 O(n)이며, 두 개의 보조 배열을 사용하므로 공간 복잡도 역시 O(n)입니다. 매 단계마다 문자열 전체를 재검사하는 O(n²) 방식보다 훨씬 효율적이어서, 입력 크기가 큰 경우에도 빠르게 동작합니다.