[Python] 백준 1419 - 등차수열의 합

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

문제

첫 항이 x이고 공차가 d인 등차수열의 첫 k개의 항은 x, x+d, x+2d, ..., x+(k-1)d이다. x와 d가 자연수인 등차수열의 첫 k개의 항의 합으로 나타낼 수 있는 수 중에서, l 이상이고 r 이하인 수가 몇 개인지 구하는 프로그램을 작성하시오.

입력

자연수 l, r, k가 순서대로 한 줄에 하나씩 주어진다. (1 ≤ l ≤ r ≤ 1,000,000,000, 2 ≤ k ≤ 5)

출력

첫째 줄에 조건을 만족하는 수의 개수를 출력한다.

 

---

 

풀이

등차수열의 합 공식은 무엇인가?

첫 항이 x이고 공차가 d라면, 등차수열의 k개의 합은

 

$$\frac{k(2x+(k-1)d)}{2}$$

 

로 나타낼 수 있다.

여기서 $l$ 이상이고 $r$이하인 것들 중에 저 값이 나올 수 있는 경우를 모두 따져보아야 할 것이다.

 

즉,

 

$$l \leq \frac{k(2x+(k-1)d)}{2} \leq r$$

 

$$\frac{2l}{k}\leq 2x+(k-1)d \leq \frac{2r}{k}$$

 

 

가장 먼저 눈에 들어오는 것은 2 ≤ k ≤ 5 라는 범위이다.

즉, k = 2, 3, 4, 5일 때를 따로 생각해주도록 하자.

 

$k = 2$일 때

 

$$l \leq 2x+d \leq r$$

 

문제 조건에서 $x$, $d$는 모두 자연수라고 했으므로, $x=1$일때를 고정시켜놓고 생각해보면 3 이상의 모든 자연수는 전부 가능하다.

 

그래서 다음을 출력해주면 된다.

 

if k == 2:
    print(max(0, r - max(l, 3) + 1))

 

(혹시 음수가 될 수도 있으니 매 코드마다 max함수로 0과 비교해주는 것을 잊지 말자)

 

$k = 3$일 때

$$\frac{l}{3}\leq x+d \leq \frac{r}{3}$$

 

math 라이브러리에서 ceil, floor 함수를 이용해서 $\frac{l}{3}$, $\frac{r}{3}$의 값을 정수로 만들어주고 시작하자.

이 경우는 2 이상의 모든 경우가 가능할 것이다.

 

elif k == 3:
    l = ceil(l/3)
    r = floor(r/3)
    print(max(0, r - max(l, 2) + 1))

 

$k = 5$일 때$(k = 4$일때는 잠시 건너뜀)

$$\frac{l}{5}\leq x+2d \leq \frac{r}{5}$$

 

이 경우는 $k=3$일때와 똑같은 방법으로, 3 이상의 모든 경우가 가능하다.

 

else:
    l = ceil(l / 5)
    r = floor(r / 5)
    print(max(0, r - max(l, 3) + 1))

 

$k = 4$일 때

이 경우가 살짝 문제가 된다. 식을 살펴보자.

 

$$\frac{l}{2}\leq 2x+3d \leq \frac{r}{2}$$

 

직접 x와 d에 수를 1부터 수를 대입해보면서 규칙을 찾아보자.

 

그렇다. 1, 2, 3, 4, 6 빼고는 모두 다 되는 것이다!

 

elif k == 4:
    l = ceil(l / 2)
    r = floor(r / 2)
    tmp = r - l + 1
    if l <= 1 <= r:
        tmp -= 1
    if l <= 2 <= r:
        tmp -= 1
    if l <= 3 <= r:
        tmp -= 1
    if l <= 4 <= r:
        tmp -= 1
    if l <= 6 <= r:
        tmp -= 1
    print(max(0, tmp))

 

 

전체 코드는 다음과 같다.

 

from math import *
l = int(input())
r = int(input())
k = int(input())

#나올 수 있는 값은 k*(2*x + (k-1)*d)//2

if k == 2:
    print(max(0, r - max(l, 3) + 1))
elif k == 3:
    l = ceil(l/3)
    r = floor(r/3)
    print(max(0, r - max(l, 2) + 1))
elif k == 4:
    l = ceil(l / 2)
    r = floor(r / 2)
    tmp = r - l + 1
    if l <= 1 <= r:
        tmp -= 1
    if l <= 2 <= r:
        tmp -= 1
    if l <= 3 <= r:
        tmp -= 1
    if l <= 4 <= r:
        tmp -= 1
    if l <= 6 <= r:
        tmp -= 1
    print(max(0, tmp))
else:
    l = ceil(l / 5)
    r = floor(r / 5)
    print(max(0, r - max(l, 3) + 1))