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

Python으로 쿼리 범위 내 서로 다른 부분 문자열 개수 구하기 (접미사 배열 + 카사이 알고리즘)

문제 소개

길이가 n인 문자열 s와 쿼리 리스트 Q가 주어졌다고 가정해 보겠습니다. 각 쿼리 Q[i]는 한 쌍의 값 (l, r)을 담고 있으며, 각 쿼리마다 문자열 s에서 인덱스 l부터 r까지(양 끝 포함) 범위에 속하는 서로 다른 부분 문자열(distinct substring)의 개수를 세어야 합니다.

예를 들어 입력이 s = "ppqpp", Q = [(1,1), (1,4), (1,1), (0,2)]라면 결과는 [1, 8, 1, 5]가 됩니다. 그 이유는 다음과 같습니다.

  • 쿼리 (1, 1): 해당 범위의 문자열은 "p" 하나뿐이므로 서로 다른 부분 문자열은 1개입니다.
  • 쿼리 (1, 4): 범위의 문자열은 "pqpp"이며, 부분 문자열은 'p', 'q', 'pq', 'qp', 'pp', 'pqp', 'qpp', 'pqpp'로 총 8개입니다.
  • 쿼리 (1, 1): 앞서와 마찬가지로 'p' 하나뿐이므로 1개입니다.
  • 쿼리 (0, 2): 범위의 문자열은 "ppq"이며, 서로 다른 부분 문자열은 'p', 'q', 'pp', 'pq', 'ppq'로 총 5개입니다.

해결 전략

이 문제는 접미사 배열(Suffix Array)카사이 알고리즘(Kasai's Algorithm)을 활용하면 효율적으로 풀 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  1. 각 쿼리마다 해당 범위의 부분 문자열 sub를 추출합니다.
  2. sub의 모든 접미사를 사전순으로 정렬해 접미사 배열을 만듭니다.
  3. 카사이 알고리즘으로 인접한 접미사들 간의 최장 공통 접두사(LCP) 배열을 선형 시간에 계산합니다.
  4. 서로 다른 부분 문자열의 총 개수는 전체 접미사 길이의 합 − 인접 접미사 간 LCP 값의 합으로 구할 수 있습니다.

알고리즘 상세 단계

kasai() 함수 — 문자열 s, 접미사 배열 suff, 길이 n을 받아 LCP 배열을 반환합니다.

  • lcp := 크기 n의 배열, 모든 값을 0으로 초기화
  • inv := 크기 n의 배열, 모든 값을 0으로 초기화
  • i를 0부터 n−1까지 순회하며 inv[suff[i]] := i 저장 (역 치환 배열 생성)
  • k := 0으로 초기화
  • i를 0부터 n−1까지 순회하며 다음을 수행:
    • inv[i] == n−1이면 k := 0으로 만들고 다음 반복으로 진행
    • j := suff[inv[i] + 1]
    • i + k < n 이고 j + k < n 이며 s[i+k] == s[j+k]인 동안 k를 1씩 증가
    • lcp[inv[i]] := k 저장
    • k > 0이면 k를 1 감소
  • lcp 반환

solve() 메인 로직 — 각 쿼리를 순서대로 처리합니다.

  • res := 빈 리스트 생성
  • 각 쿼리 (left, right)에 대해:
    • sub := s[left:right+1], length := right − left + 1
    • suffix := [(i, sub[i:]) for i in range(length)] 형태의 (인덱스, 접미사) 쌍 리스트 생성
    • suffix를 접미사(두 번째 요소) 기준으로 정렬
    • 정렬 결과에서 인덱스 배열 suff와 접미사 배열 suffix를 분리
    • lcp := kasai(sub, suff, length) 호출
    • count := len(suffix[0])으로 시작해, i를 0부터 length−2까지 count += len(suffix[i+1]) − lcp[i] 누적
    • count를 res에 추가
  • 모든 쿼리 처리 후 res 반환

구현 예제

아래 파이썬 코드를 통해 더 자세히 이해해 보겠습니다.

def kasai(s, suff, n):
    lcp = [0] * n
    inv = [0] * n
    for i in range(n):
        inv[suff[i]] = i
    k = 0
    for i in range(n):
        if inv[i] == n - 1:
            k = 0
            continue
        j = suff[inv[i] + 1]
        while i + k < n and j + k < n and s[i + k] == s[j + k]:
            k += 1
        lcp[inv[i]] = k
        if k > 0:
            k -= 1
    return lcp

def solve(s, Q):
    res = []
    for i in range(len(Q)):
        left, right = Q[i]
        sub = s[left:right + 1]
        length = right - left + 1

        suffix = [[i, sub[i:]] for i in range(length)]

        suffix.sort(key=lambda x: x[1])
        suff, suffix = [list(t) for t in zip(*suffix)]

        lcp = kasai(sub, suff, length)
        count = len(suffix[0])
        for i in range(length - 1):
            count += len(suffix[i + 1]) - lcp[i]

        res.append(count)
    return res

s = "pptpp"
Q = [(1, 1), (1, 4), (1, 1), (0, 2)]
print(solve(s, Q))

입력

"pptpp", [(1,1),(1,4),(1,1),(0,2)]

출력

[1, 8, 1, 5]

동작 원리 요약

정렬된 접미사 목록에서 각 접미사는 그 자체로 여러 개의 부분 문자열(접두사)을 포함합니다. 따라서 모든 접미사 길이의 합이 잠재적인 부분 문자열의 총 개수가 됩니다. 그런데 인접한 두 접미사가 lcp[i]만큼 공통 접두사를 공유한다면, 그만큼 중복되는 부분 문자열이 발생합니다. 결국 전체 접미사 길이의 합에서 인접 접미사 간 LCP의 합을 빼면 서로 다른 부분 문자열의 개수가 정확히 계산됩니다.

시간 복잡도

각 쿼리마다 접미사 정렬에 O(m log m)(m은 쿼리 범위 길이, 비교당 최대 O(m)), 카사이 알고리즘에 O(m)이 소요됩니다. 따라서 총 시간 복잡도는 쿼리당 대략 O(m² log m) 수준이며, 문자열이 짧거나 쿼리 수가 많지 않은 경우에 실용적으로 동작합니다.