[Python 3] BOJ - 3078 "좋은 친구"

 # 문제 링크

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

3078번: 좋은 친구

 

 # 풀이

 개인적으로 이 문제의 핵심은 upper bound의 활용이라고 생각한다. 문제에서 제시한 좋은 친구의 조건은 점수의 차이가 K이하이면서 같은 이름을 갖는 친구이다. 따라서 우선적으로 이름이 같은 사람끼리 성적을 비교해야 한다. 따라서, 이름의 길이를 key, 그 이름의 길이를 갖는 사람의 성적 리스트를 value로 갖는 dictionary가 필요하다.

<그림 1. dictionary의 예시>

 이 때 각각의 value는 정렬된 상태여야 한다. 따라서, 각각의 value list의 원소에서 자신보다 k보다 작거나 같은 사람은 좋은 친구가 될 수 있다. 따라서, k보다 큰 사람은 좋은 친구가 될 수 없다. 따라서, 현재 value list의 인덱스에서 k만큼 차이나는 upper bound의 인덱스를 찾은 후, 빼주고 다시 1을 빼주면 좋은 친구의 중복없는 쌍을 구할 수 있다. 1을 다시 빼주는 이유는 자기 자신이 포함되어 있기 때문이다.

 

 # 코드

import sys, bisect

# 입력부
n, k = map(int, sys.stdin.readline().split())

# arr : 이름의 길이를 key, 이름의 길이에 해당하는 사람들의 성적 리스트를 value로 갖는 dict
arr = dict()
for i in range(n):
    temp = sys.stdin.readline().rstrip()
    size = len(temp)
    # arr에 이름 사이즈와 성적을 추가, 이 때 자연스럽게 정렬
    if size not in arr:
        arr[size] = [i + 1]
    else:
        arr[size].append(i + 1)
ans = 0

for i in arr:
    for j in range(len(arr[i])):
        # bisect로 upper_bound의 index를 구하고 빼준 후, 1을 다시 뺌
        ans += (bisect.bisect_right(arr[i], arr[i][j] + k) - j - 1)
        
# 정답 출력
print(ans)