[Python/파이썬] 백준 20437번 문자열 게임 2

20437번: 문자열 게임 2

 

문제

작년에 이어 새로운 문자열 게임이 있다. 게임의 진행 방식은 아래와 같다.

1. 알파벳 소문자로 이루어진 문자열 W가 주어진다.
2. 양의 정수 K가 주어진다.
3. 어떤 문자를 정확히 K개를 포함하는 가장 짧은 연속 문자열의 길이를 구한다.
4. 어떤 문자를 정확히 K개를 포함하고, 문자열의 첫 번째와 마지막 글자가 해당 문자로 같은 가장 긴 연속 문자열의 길이를 구한다.

위와 같은 방식으로 게임을 T회 진행한다.

입력

문자열 게임의 수 T가 주어진다. (1 ≤ T ≤ 100)

다음 줄부터 2개의 줄 동안 문자열 W와 정수 K가 주어진다. (1 ≤ K ≤ |W| ≤ 10,000) 

출력

T개의 줄 동안 문자열 게임의 3번과 4번에서 구한 연속 문자열의 길이를 공백을 사이에 두고 출력한다.

만약 만족하는 연속 문자열이 없을 시 -1을 출력한다.

코드

처음 풀이한 코드 => 시간 초과

from collections import defaultdict

for _ in range(int(input())):
    w = input()
    k = int(input())
    
    # 3번
    start, end = 0, 0
    count = defaultdict(int)
    min_len = 10001
    
    while end < len(w):
        count[w[end]] += 1
        if count[w[end]] == k:  # 어떤 문자를 정확히 K개 포함
            while count[w[end]] == k and w[start] != w[end]:
                count[w[start]] -= 1
                start += 1
            min_len = min(min_len, end-start+1)
            count[w[start]] -= 1
            start += 1
        end += 1
    # 만족하는 연속 문자열이 없는 경우
    if min_len == 10001:
        print(-1)
    else:
        # 4번
        max_len = -1
        flag = False
        for length in range(len(w), -1, -1):
            if flag:
                break
            for i in range(len(w)-length):
                if w[i] == w[i+length-1] and w[i:i+length].count(w[i]) == k:
                    max_len = length
                    flag = True
                    break
        print(min_len, max_len)

예제는 통과하였으나 제출 시 시간 초과가 발생한 코드이다.

 

3. 어떤 문자를 정확히 K개를 포함하는 가장 짧은 연속 문자열의 길이를 구한다.

=> 투포인터를 이용하여 풀이하려고 하였다. defaultdict를 이용하여 문자열에 등장한 문자들의 개수를 저장해준다. end 지점의 문자의 개수가 k개가 된다면 그 문자열은 end 지점의 문자를 k개 포함하고 있는 문자열인 것이다. 이 때 우리가 구하려는 가장 짧은 연속 문자열은 start의 문자가 end와 같을 때의 문자열일 것이다. 그래서 그 지점을 찾을 때까지 while문을 돈다.

4. 어떤 문자를 정확히 K개를 포함하고, 문자열의 첫 번째와 마지막 글자가 해당 문자로 같은 가장 긴 연속 문자열의 길이를 구한다.

=> 가장 긴 길이의 문자열부터 살펴보며 양쪽 끝의 글자가 같고 그 글자가 현재 보고 있는 문자열에 포함된 개수가 k개인 것을 찾으면 반복문을 탈출하도록 하였다.

아마 여기서 시간 초과가 발생한 것 같다. count 함수의 시간복잡도는 O(n)으로 이 코드는 최악의 경우 시간복잡도가 O(n^3)이다. 

 

정답 코드

from collections import defaultdict

def str_game(w, k):
    over_k = defaultdict(list)
    for i in range(len(w)):
        # w에 k개 이상 포함된 문자를 over_k 딕셔너리에 {문자 : 위치리스트} 형태로 저장
        if w.count(w[i]) >= k:
            over_k[w[i]].append(i)
    if len(over_k) == 0:
        return [-1]
    min_len = 10001
    max_len = 0
    for pos_list in over_k.values():
        for j in range(len(pos_list)-k+1):
            # 해당 문자를 k개만큼 포함한 문자열의 길이 : pos_list[j+k-1]-pos_list[j]+1
            min_len = min(min_len, pos_list[j+k-1]-pos_list[j]+1)
            max_len = max(max_len, pos_list[j+k-1]-pos_list[j]+1)
    return [min_len, max_len]

for _ in range(int(input())):
    w = input()
    k = int(input())
    
    print(*str_game(w, k))

• str_game()
  • over_k : 입력으로 주어진 문자열에서 k개 이상 포함된 문자들의 위치만 저장하는 defaultdict. 
  • 만약 over_k의 길이가 0이라면 문자열에 k개 이상 포함된 문자가 없는 것이므로 문제의 조건을 만족하는 연속 문자열을 만들 수 없다. 따라서 -1을 리턴.
  • over_k의 value 값에는 key값 문자의 인덱스들이 리스트로 저장되어 있다. 해당 문자를 k개만큼만 포함한 문자열의 길이는 이 리스트의 값들을 갖고 구할 수 있다. 아래의 식이 문자를 k개만큼 포함한 연속 문자열의 길이를 구하는 식이다.

pos_list[j+k-1]-pos_list[j]+1