세 개의 값 n, total, k가 주어졌다고 가정해 봅시다. 이때 다음 조건을 만족하는 크기 n짜리 리스트를 생각합니다.
- 리스트의 모든 요소의 합이 total과 같다.
- 리스트에서 인접한 두 요소의 절대 차이는 최대 1이다.
이러한 조건을 만족하는 리스트 중에서 인덱스 k 위치에 놓일 수 있는 최댓값을 구하는 것이 문제입니다.
예를 들어 n = 5, total = 15, k = 3이 주어지면 출력은 4가 됩니다. 그 이유는 [3, 2, 3, 4, 3] 같은 리스트가 조건을 만족하며, 이 리스트의 인덱스 3에 있는 값이 4로 가능한 경우 중 가장 크기 때문입니다.
풀이 접근 방법
이 문제는 후보 값을 하나씩 키워가며 검증하는 방식으로 해결할 수 있습니다. 어떤 값 x를 인덱스 k에 배치하면, 인접 요소 차이가 최대 1이라는 조건 때문에 왼쪽과 오른쪽으로 갈수록 값이 최대 1씩 감소하는 형태가 됩니다.
구체적인 단계는 다음과 같습니다.
- x := 0 으로 초기화합니다.
- 아래 과정을 반복합니다.
- a := k + 1 (인덱스 0부터 k까지의 요소 개수)
- s := (x + x - a + 1) * floor(a / 2) — 왼쪽 부분을 등차수열의 합 공식으로 계산
- a := n - k (인덱스 k부터 n-1까지의 요소 개수)
- s := s + (x + x - a + 1) * floor(a / 2) — 오른쪽 부분의 합을 더함
- s := s - x — 인덱스 k 위치의 값이 두 번 더해졌으므로 한 번 제거
- 만약 s > total이면 루프를 빠져나갑니다.
- x := x + 1
- x - 1을 반환합니다.
예제 코드
아래 구현을 통해 더 잘 이해해 보겠습니다.
def solve(n, total, k):
x = 0
while 1:
a = k + 1
s = (x + x - a + 1) * a // 2
a = n - k
s += (x + x - a + 1) * a // 2
s -= x
if s > total:
break
x += 1
return x - 1
n = 5
total = 15
k = 3
print(solve(n, total, k))입력
5, 15, 3
출력
4
동작 원리 정리
위 코드는 인덱스 k에 값 x를 놓았을 때 전체 리스트의 합이 얼마가 되는지를 등차수열의 합 공식으로 빠르게 계산합니다. 왼쪽 구간(0 ~ k)과 오른쪽 구간(k ~ n-1)의 합을 각각 구한 뒤 중복된 x를 빼주면 전체 합이 됩니다. 이 합이 total을 초과하는 순간의 직전 값, 즉 x - 1이 인덱스 k에 놓일 수 있는 최댓값입니다. 시간 복잡도는 O(total) 수준으로, 각 후보 값마다 상수 시간의 계산만 필요하기 때문에 효율적입니다.