[Python/파이썬] 백준 1241번 머리 톡톡

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

코드

import sys
from collections import defaultdict
input = sys.stdin.readline

n = int(input())
students = []
num_cnt = defaultdict(int)
for _ in range(n):
    input_value = int(input())
    students.append(input_value)
    num_cnt[input_value] += 1

maximum_value = max(students)

count = [0]*(maximum_value+1)
for k, v in sorted(num_cnt.items()):
    count[k] += v-1
    for i in range(k*2, maximum_value+1, k):
        count[i] += v

answer = []
for i in range(n):
    answer.append(count[students[i]])

print(*answer, sep="\n")

 

문제를 잘 읽다보면 "배수" 이 부분에서 에라토스테네스의 체를 사용해야겠다는 느낌이 온다.

 

근데 입력 값이 오름차순이 아니기도 하고, 출력할 때도 입력받은 순으로 해야 하기 때문에 입력받은 수와 그 개수를 defaultdict에 저장해놓고, 이 딕셔너리로 에라토스테네스의 체 알고리즘을 수행했다.

 

나보다 크거나 같은 배수에 표시를 해줘야 한다. 그렇기 때문에 각 숫자 k와 그 개수 v에 대해, k의 배수들(k, 2k, 3k, ...)에 이 숫자가 등장한 횟수 v만큼 카운트를 추가해 준다.

다만, 자기 자신은 "자신보다 큰 수"가 아니므로, 자신의 개수에서 1을 뺀 값(v-1)만 더해준다. 이는 자기 자신과 같은 숫자가 여러 번 등장할 경우, 그 중 하나를 제외한 나머지만 카운트해야 하기 때문이다.

 

count 배열은 각 인덱스 i에 대해, i보다 작으면서 i를 나눌 수 있는 숫자들의 등장 횟수 합을 저장하게 된다. 마지막으로 입력받은 순서대로 각 학생 번호에 해당하는 count 값을 출력한다.