[Algorithm🧬] 백준 3078 : 좋은 친구

문제

MEMORY LIMIT EXCEED

솔직히 TIME LIMT EXCEED 날 줄 알았는데 메모리 초과 나서 좀 놀랐지만.. 어쨌든 다시 풀어야한다는 사실에 너무 머리가 안돌아가서 절망했으나

import sys
from itertools import combinations as cb

readl = sys.stdin.readline

n, k = map(int, readl().split())
students = [[i, name, len(name)] for i, name in enumerate([readl().rstrip() for _ in range(n)])]

# print(students)

bestfriend = set()

for i in range(n - k):
    arr = []

    # 후보군을 넣고
    candidates = students[i:i+k+1]
    candidates.sort(key=lambda x:x[2])

    # print(candidates)

    length = candidates[0][2]
    local = []
    j = 0
    while j < k + 1:
        if length == candidates[j][2]:
            local.append(candidates[j][1])
            j += 1
        else:
            arr.append(local)
            local = []
            length += 1

    arr.append(local)
    local = []
    length += 1

    # print(arr)

    for a in arr:
        for c in cb(a, 2):
            bestfriend.add(c)


print(len(bestfriend))

큐를 이용해서 푼 답!!!

다음거 풀려고 보다가 갑자기 일단 처음에 들어가는 것만 nC2 고 다음건 + (n - 1) 하면 되는거 아닌가? 하고 불현듯이 떠오르고,

나룻배 문제에서 큐를 left right 로 관리 안하고 인덱스 붙여서 관리한게 생각나면서, 아래와 같은 풀이를 해버림... 나 좀 쩌는걸지도

import sys
from collections import deque

readl = sys.stdin.readline

def nC2(n):
    return int(n * (n - 1) / 2)


n, k = map(int, readl().split())
students = [[i, name, len(name)] for i, name in enumerate([readl().rstrip() for _ in range(n)])]

# 큐를 이름 길이별로 운영
queue = [deque() for _ in range(21)] # 2글자부터 20글자
sum = 0

# 처음 K개 넣고 시작.
for i in range(k + 1):
    _, name, length = students[i]
    queue[length].append(name)
# 처음 K 개에 대한 베프 쌍 개수
# 처음에는 k range 안에 있고 이름 길이 같으면 전부 다 친구가 될 수 있음 nC2
for i in range(2, 21):
    if queue[i]:
        sum += nC2(len(queue[i]))

# print(queue, sum)

# 그리고 나서는 맨앞의거 하나 없애고, 맨 뒤에거 하나 넣어주는 식으로 
# 대신 이미 큐에 들어간 애들은 베프가 되었으니 새로운 애에 대해서 나머지 친구들과 베프가 되면 됨.
for i in range(k + 1, n):
    _, name, length = students[i]
    queue[students[i - k - 1][2]].popleft()
    queue[length].append(name)

    sum += (len(queue[length]) - 1)

    # print(queue, sum)

print(sum)