소수 찾기

출처 | 소수 찾기[백준 1978]

문제

무한히 큰 배열에 다음과 같이 분수들이 적혀있다.

이와 같이 나열된 분수들을 1/1 → 1/2 → 2/1 → 3/1 → 2/2 → … 과 같은 지그재그 순서로 차례대로 1번, 2번, 3번, 4번, 5번, … 분수라고 하자.

X가 주어졌을 때, X번째 분수를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 X(1 ≤ X ≤ 10,000,000)가 주어진다.

출력

첫째 줄에 분수를 출력한다.

---

풀이

import sys
import math

N = int(input())
M = list(map(int, sys.stdin.readline().split()))
count = 0

def primenumber(M):
  for i in range(2, int(math.sqrt(M)) + 1):
    if M % i == 0:
      return 0
  return 1

for i in range(N): 
  count += primenumber(M[i])
  if M[i] == 1 and count > 0:
    count -= 1
print(count)