3078. 좋은 친구 (Python)

 

3078번: 좋은 친구

 

풀이

 슬라이딩 윈도우 개념을 사용하여 쉽게 해결할 수 있는 문제다. 이름의 글자수에 대한 리스트(lengths)를 하나 만들어 n글자의 이름이 현재 범위에서 총 몇 개가 있는지를 관리한다.

 

 위와 같은 예제를 통해 이해해보자. 먼저 처음 CYNTHIA 이름의 길이(now)인 7을 기억한다. 그 후 K가 3이므로 LLOYD ~ KEVIN의 범위에서 이름들의 길이를 가져와 lengths에 저장한다. 그러면 lengths[5] = 2 (LLOYD, KEVIN), lengths[6] = 1 (STEVIE)이 된다. 그 후 lengths[now]를 통해 현재 기준이 되는 이름의 길이에 해당하는 이름이 몇 개가 있는지를 기록한다.

 그 후 기준을 LLOYD로 옮기는데 이때, 범위는 STEVIE ~ MALCOLM으로 옮겨간다. 이 때 이전 범위와 현재 범위에서 STEVIE와 KEVIN은 중복되기 때문에 다시 기록해줄 필요가 없다. 범위를 수정할 때는 이전 범위에서 제거되는 LLOYD의 길이를 lengths에서 제거해주고 새롭게 범위에 추가되는 MALCOLM에 대한 길이만 lengths에 추가해주면 된다. 그 후 다시 한 번 lengths[now]를 통해서 이름이 몇 개가 있는지 체크한다.

 

 위와 같은 과정을 N-1까지 반복해주면 된다.

 

코드

if __name__ == '__main__':
    N, K = map(int, input().split())  # N: 학생의 수, K: 자신과 반 등수의 차이
    friends = []  # 학생의 이름
    lengths = [0 for _ in range(21)]  # 살펴볼 범위에서 학생들의 이름의 길이의 수 카운트

    for _ in range(N):
        friends.append(input())

    now = len(friends[0])  # 먼저 처음 학생이 기준이 된다.

    for i in range(1, 1+K):  # 처음 학생으로부터 K만큼 떨어진 학생까지의 이름의 길이를 가져와 카운트한다.
        lengths[len(friends[i])] += 1

    res = lengths[now]  # 기준이 되는 학생과 친구가 될 수 있는 친구의 수를 저장한다.

    for i in range(1, N-1):  # 이제 다음 학생부터 차례대로 범위를 바꿔가며 친구의 수를 파악한다.
        now = len(friends[i])
        if i+K < N:
            lengths[len(friends[i+K])] += 1  # 범위에 새롭게 추가될 학생이 있다면 추가한다.

        lengths[now] -= 1  # 이제 기준이 될 학생은 범위에서 제거한다
        res += lengths[now]  # 기준이 되는 학생과 친구가 될 수 있는 친구의 수 저장

    print(res)