[Python][백준/BOJ] 3078번 : 좋은 친구

https://www.acmicpc.net/problem/3078

3078번: 좋은 친구

 

내 코드

import sys
input = sys.stdin.readline

n, k = map(int, input().split())
window = {i:0 for i in range(2,21)} #2는 몇개, 3은 몇 개, 4는 몇 개.. 이름 길이별로 개수 세서 저장함
student=[0]*n
ans =0      

    
for i in range(n):
    f = str(input().rstrip())
    student[i] = len(f)
    if i > k: #윈도우 이동 시
        window[student[i-k-1]] -= 1
    ans += window[len(f)] # 먼저 +1을 해주면 안되는 이유는 문제 요구하는 것은 2인 이상의 쌍이기 때문
    #해당 len(f)는 현재 기준으로 두고 있는 학생을 의미함, 현재 학생의 길이가 window 안에 몇 개 있는지를 ans에 저장함(아직 +1을 안해줬으니까 현재 학생을 제외한 개수)
    window[len(f)] +=1
    
print(ans)

 

코드 리뷰

슬라이딩 윈도우

이번 문제도 푸는데 꽤 오랜 시간이 걸렸다. 슬라이딩 윈도우 문제이지만, 두 번 생각해서 풀어야 하는 문제같이 느껴졌다.

슬라이딩 윈도우는 부분 합의 문제이기 때문에 부분 합을 요구한다면 가볍게 접근할 수 있는 문제같다. 하지만 만약 슬라이딩 윈도우의 문제이나, 합을 요구하지 않고 이외의 문자열이나 혹은 문자의 길이 등과 같은 조건이 주어진다면 갑자기 난이도가 높아지는 것 같다고 생각했다.

그리고 길이나 문자열이 주어진다면, dictionary를 사용하는 것이 하나의 요령인 것 같다. 시간복잡도도 짧을 뿐만 아니라 구분해서 윈도우 구간에 들어간다면 넣고 윈도우 밖으로 벗어났다면 빼는 것이 슬라이딩 윈도우에 딱 잘 어울리기 때문이다. 

만약 순서가 중요하다면 deque을 사용할 수도 있다. 사용하는 도구는 여러가지이지만, 확실히 자주 사용되는 적절한 도구들이 몇 개가 있는 것 같다.

 

이 문제는 문자열들이 등장하지만, 사실 가장 중요한 것은 문자의 길이이다. 그래서 아예 입력받을 때부터 len(input)으로 받아버려도 된다.

    f = str(input().rstrip())
    student[i] = len(f)

이렇게 입력받은 문자열의 길이를 student라는 배열 안에 저장했다. 순서대로 리스트 안에 차곡차곡 길이를 저장한 것이다.

    if i > k: #윈도우 이동 시
        window[student[i-k-1]] -= 1

이 구문을 하나씩 찬찬히 뜯어보자. 먼저 이 문제를 풀 때는 if i > k라는 조건을 사용했다. k는 윈도우의 크기이며, i > k 가 의미하는 것은 윈도우가 옆으로 이동한 이후를 의미한다.

student[i-k-1]

 

윈도우의 왼쪽과 오른쪽이 있다면, 계속 윈도우는 오른쪽으로 한 칸씩 이동할 것이다. 그러면 윈도우의 왼쪽에 해당하는 부분은 하나씩 윈도우에서 벗어난다. 이건 윈도우의 왼쪽에 해당하며, 윈도우의 구간에서 벗어난 부분이다.

window[student[i-k-1]] -= 1

나는 슬라이딩 윈도우를 풀 때 무조건 window에 해당하는 부분을 window라고 이름짓는다. 그래야 코드가 복잡해지고, 문제가 어려워지더라도 알고리즘의 개념을 헤매지 않기 때문이다.

window = {i:0 for i in range(2,21)} #2는 몇개, 3은 몇 개, 4는 몇 개.. 이름 길이별로 개수 세서 저장함

이번 window는 dictionary 형태로 만들어져 있다. 이렇게 생겼는데, 이는 2부터 20까지 입력될 수 있는 이름들의 개수를 보관하는 dict이다. 그래서 저 구문은 결국 window에서 방금 막 벗어나서 아웃된(이제 필요없는) 녀석의 이름 길이를 window에서 찾아내서 -1을 하겠다는 말이다.

    ans += window[len(f)] # 먼저 +1을 해주면 안되는 이유는 문제 요구하는 것은 2인 이상의 쌍이기 때문
    #해당 len(f)는 현재 기준으로 두고 있는 학생을 의미함, 현재 학생의 길이가 window 안에 몇 개 있는지를 ans에 저장함(아직 +1을 안해줬으니까 현재 학생을 제외한 개수)
    window[len(f)] +=1

이렇게까지 하면 이 복잡한 문제를 해결할 수 있는 코드가 끝난다. ans에는 현재 학생의 길이와 일치하는 개수를 window에서 찾아내어 더해줄 것이다.(문제에서 요구하는 것 그 자체)

그리고 window[len(f)]+=1 을 이 이후에 해주는 이유는 문제 자체는 친구 쌍을 구하는 것이기 때문이다. 2명이 모였을 때 1쌍이 이루어지므로, 자기 자신을 카운팅하는 것은 의미없다.

이렇게 n번 반복하면 해결된다.

 

슬라이딩 윈도우를 풀어보면서 다양한 문제들을 마주치고 있는데, 개념 자체는 쉬운 윈도우 알고리즘이지만 응용하면 언제든지 더 어려워질 수 있는 알고리즘인 것 같다. 카카오에서도 자주 출제되고 있다고 하니 주의해서 많이 풀어보는 것이 좋을 듯 하다.