[BOJ] - python 스터디 그룹 (정렬 + 그리디 + 투포인터)

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

14572번: 스터디 그룹

문제는 이러하다, 보자마자 오,, 쉽겠는데? 하고 풀었다가 거진 5시간을 버린 문제이기도 해서 정리해본다.

 

우선 문제를 보자마자 $E$ = (그룹 내의 학생들이 아는 모든 알고리즘의 수 - 그룹 내의 모든 학생들이 아는 알고리즘의 수) * 그룹의 수라고 하길래 그룹의 수를 가장 올리면 그 나머지 항은 자연스럽게 올가갈 것이라 생각했다. 그래서 그리디를 생각해 냈고, 문제를 곱씹어 보던중, 이건 투포인터를 쓰면 간단히 해결할 수 있을거라고 했다.. 그럴거라고 믿었지 ㅎㅎ

import sys
input = sys.stdin.readline
 
def calcE(s, e, p_set):
    interval = e - s + 1
    f = len(list(p_set.keys()))
    ff = len(list(filter(lambda v: v >= interval, p_set.values())))
    return (f - ff) * interval
    
N, K, D = map(int, input().split())
 
S = []
for i in range(N):
    M, d = map(int, input().split())
    S.append((M, d, list(map(int, input().split()))))
    
S = sorted(S, key=lambda x: x[1])
ans = -1; s = 0; e = 0
p_set = {}
 
while(1):
    if (S[e][1] - S[s][1] <= D):
        for t in S[e][2]:
            if t in p_set.keys(): p_set[t] += 1
            else: p_set[t] = 1
        ans = max(ans, calcE(s, e, p_set)) # update
        e += 1
    else:
        for t in S[s][2]:
            p_set[t] -= 1
            if p_set[t] == 0: del p_set[t]
        s += 1
    if e >= N: break
print(ans)

코드 자체는 매우 간단하게 짰다. 우선 모든 학생들을 입력받고, 학생들의 실력을 기준으로 오름차순으로 S배열을 정렬해주었다. 그리고 투포인터를 적용해주기 위해 $s=0, e=0$으로 설정해주었고, p_set은 학생들의 알고 있는 알고리즘의 Union, Intersect를 저장해주고 빼주기 위해 만든 딕셔너리이다.

 

투포인터의 분기로 s, e가 가리키고 있는 수준의 차이가 D보다 작거나 같으면, 알고리즘들을 p_set에 추가하고 ans를 최댓값으로 업데이트 해주는 과정을 진행했다. 아니라면 s가 가리키고 있는 정보를 p_set에서 빼주었다. 이 과정은 이전에 mo's 알고리즘을 배울때 짜보았기 때문에 쉽게 할 수 있었다.

 

그리고 e >= N일때 break해주면 투포인터 알고리즘 구현이 끝난다. 이 알고리즘의 시간 복잡도는 투포인터 O(N)에서 그 안의 for문을 생각해준다면 문제는 최대  1~30개이므로 이를 K라고 했을 때, O(NK)가 될 것이여서 TLE가 나지 않는다.

 

나는 매우 틀렸었는데, 그 이유는 투포인터를 오랜만에 짜봐서 그런지, s, e를 어떠헤 설정해주었었는지 까먹어서 내 맘대로 설정하다 그냥 계속 틀렸던 것이였다.. 앞으로는 기억이 안나면 정확한 예시와 정의를 보고, 문제 조건도 꼼꼼히 보고 PS해야겠다!!