Coding 공부

import sys
from collections import defaultdict
def input():
    return sys.stdin.readline().rstrip()

N,M = map(int,input().split())

arr = list(map(int,input().split()))


prefix_sum = [0] + arr[:]
result = 0
num_dict = defaultdict(int)
for i in range(1,N+1):
    prefix_sum[i] += prefix_sum[i-1]
    calc = prefix_sum[i] - M*i

    result += num_dict[calc]
    num_dict[calc] += 1


print(result+num_dict[0])

몇달 전에 푼 문제라 정확히 복기는 못하지만, 현재 생각 나는대로 복기를 할려고 한다.

 이 문제는 누적합을 이용해 구간의 평균합의 개수를 구하는 것이다.

현재까지의 누적합에서 지금까지의 평균의 총합을 빼주게 된다면, 

현재 위치에서 평균이 되기 위해 사라져야할 값이 보인다.

이 점을 이용해 푸는 문제이다.

구해야하는 평균이 K일때

즉 현재의 누적합 - 구간*K 을 했을때 나오는 값을 calc이라 했을 때,

이 calc 만큼만 사라지면 현재 구간 앞에서 얻을 수 있는 평균이 K인 구간의 개수가 되는 것이다.

이 문제는 누적합에서 구간의 평균합을 뺀 숫자와 일치하는 개수를 계속 누적해서 더해주면 풀 수 있는 문제이다.

그리고 마지막에 0을 더해준 이유는 현재 구간의 앞에서 나올 수 있는것을 더해준 것이니,

전체 구간에서 더해줄 필요가 있어서, 마지막에 더해주는 것이다.

6달전에 풀었던 것을 다시 복기할려니 정확하게 설명은 못했지만, 위와 같은 방법을 활용하면 된다.

import sys
input = sys.stdin.readline

N = int(input())


problem_dict = {}
re1_dict = {}

for _ in range(N):
    pb_num,l_num = map(int,input().split())
    re1_dict[(l_num,pb_num)] = 1
    problem_dict[pb_num] = l_num


M = int(input())

for _ in range(M):
    command,*arg = input().split()

    if command == 'add':
        pb_num,l_num = map(int,arg)
        problem_dict[pb_num] = l_num
        re1_dict[(l_num,pb_num)] = 1
    elif command == 'recommend':
        flag = int(arg[0])
        if flag > 0:
            keys = max(re1_dict.keys())
        else:
            keys = min(re1_dict.keys())
        print(keys[1])
    else:
        pb_num = int(arg[0])
        l_num = problem_dict[pb_num]
        del problem_dict[pb_num]
        del re1_dict[(l_num,pb_num)]

 이 문제를 풀때 사실 boj.kr/21944 번을 먼저 풀었기 때문에 그 때 시도했던 후자 방식을 먼저 적용시켜보았다.

제한 시간내에 동작은 하지만 위 풀이는 추천하지 않는다.

문제를 푸는 아이디어는 다음과 같다. 어차피 우리는 recommend를 해야하고, 그때의 값들을 관리하는 것이다.

그러면 이 문제를 풀때 dictinory를 이용하기로 했다.

(난이도, 문제번호)를 key 값으로 하는 dictionary를 만들어준 뒤 

recommend를 해줄때 1일때에는 최대값

-1일때는 최소값을 구해서 출력해주는 식으로 했다.

max,min이 알아서 앞에서부터 비교를 해주니 가능한 방식이다.

그리고 문제를 풀었을 때에는

problem_dict 이라는 dictionary에 문제번호와 난이도를 매핑시켜준 딕셔너리가 있다.

그러면 problem_dict에서 solved 된 문제번호를 통해 난이도를 가져오고,

re1_dict의 값을 제거해주면 된다.

import sys
import heapq
input = sys.stdin.readline

def recommend(flag,heap):
    flag = -flag
    while heap and (heap[0][1]*flag not in problem_dict.keys() or problem_dict[heap[0][1]*flag] != heap[0][0]*flag):
        heapq.heappop(heap)
    result =heap[0]
    result = result[1]*flag
    return result

N = int(input())
max_heap = []
min_heap = []
problem_dict = {}
for _ in range(N):
    pb_num,l_num = map(int,input().split())
    max_heap.append((-l_num,-pb_num))
    min_heap.append((l_num,pb_num))

    problem_dict[pb_num] = l_num
heapq.heapify(max_heap)
heapq.heapify(min_heap)
M = int(input())

for _ in range(M):
    command,*arg = input().split()
    if command == 'add':
        pb_num,l_num = map(int,arg)
        heapq.heappush(max_heap,(-l_num,-pb_num))
        heapq.heappush(min_heap,(l_num,pb_num))
        problem_dict[pb_num] = l_num
    elif command == 'solved':
        pb_num = int(arg[0])
        del problem_dict[pb_num]
    else:
        flag = int(arg[0])
        if flag > 0:
            print(recommend(flag,max_heap))
        else:
            print(recommend(flag,min_heap))

이게 정석적인 풀이이다.

min_heap 과 max_heap을 이용해서 풀었다.

여기서 주의해야할점은 problem_dict을 활용해 사라진 문제들을 찾아서, 제거를 해주어야한다.

while문을 통해 없는 문제들은 전부 없애주고, while문을 탈출했을때 heap[0]에 있는 값을 return 해주면 된다.