[Python] 최고의 집합

최고의 집합

• 문제 출처: 프로그래머스
• 난이도: level 3
• 참고: 없음
• input size: $1 \le n \le 10^4$, $1 \le s \le 10^8$
• (주관) 자료구조: 없음
• (주관) 알고리즘: 수학
• (주관) 예상 시간복잡도: $O(n)$

합이 $s$인 $n$개의 자연수 집합을 구하고, 그 중 원소의 합이 최대인 집합을 오름차순으로 정렬한다.

 

STEP 1. 관찰과 직관

프로그래머스에 주어진 예시가 $n=2$라서 와닿지는 않는다.

왜냐면 $n=2$인 경우에 너무 쉽게 찾을 수 있다.

$a_1 = x, x_2 = s-x$라 하면 곱 $f(x) = x(s-x)$이고 이차함수의 최댓값은 $x=s/2$인 지점이다.

물론 정수이므로 $a_1 = \lfloor s/2 \rfloor$, $a_2 = s-a_1$으로 간단히 찾을 수 있다.

 

$n=3$인 예를 통해 규칙을 찾아보자.

$n=3, s=9$인 경우, {1, 1, 7}, {1, 2, 6}, {1, 3, 5}, {1, 4, 4}, {2, 2, 5}, {2, 3, 4}, {3, 3, 3}이 가능한 집합이고, 이 중 곱이 최대가 되는 집합은 {3, 3, 3}이다.

$n=3, s=8$인 경우, {1, 1, 6}, {1, 2, 5}, {1, 3, 4}, {2, 2, 4}, {2, 3, 3}이 가능한 집합이고, 이 중 곱이 최대인 집합은 {2, 3, 3}이다.

뭔가 가능한 원소들이 균일하게(uniform)한 값을 가지면 될 듯 하다.

STEP 2. 산술-기하 평균 부등식

음이 아닌 실수 $a_1, \dots, a_n$에 대하여 다음 부등식이 성립한다.

\[ \frac{a_1 + a_2 + \cdots + a_n}{n} \ge \sqrt[n]{a_1a_2 \cdots a_n} \]

이때, 등호가 성립하는 조건은 $a_1 = a_2 = \cdots = a_n$이다.

 

따라서 $s = a_1 + a_2 + \cdots + a_n$이 고정된 값으로 주어지므로, 우리는 곱의 최댓값을 구할 수 있다.

정확히는, 곱의 최댓값을 구할 필요가 없으므로 $a_i = s/n$ 임을 알 수 있다.

STEP 3. 정수 조건

그런데 위 산술-기하 평균 부등식은 "실수"에 대해서 성립한다.

우리가 원하는 정답은 "정수"이다.

$a_i = s/n$이 정수라는 보장은 없다.

 

정수는 항상 $a = qb + r (0 \le r < q)$로 나타낼 수 있으므로 $s/n$의 나머지 $r$에 따라 달라진다.

$r=0$이면 $s/n$은 정수이므로 모든 $a_i$가 같은 값 $q = s/n$을 가진다.

만약 $r>0$이면 $r$개의 원소들은 $1$씩 더해서 $q+1$의 값을 가진다.

따라서 $n-r$개의 원소들은 $q$, $r$개의 원소들은 $q+1$을 갖는다.

 

최종 코드

최종 소스 코드는 다음과 같다.

def solution(n, s):
    q = s // n
    r = s % n
    if q == 0:
        answer = [-1]
    else:
        answer = [q] * (n-r) + [q+1] * r
    return answer

 

$q=0$인 경우, 애초에 $n>s$이므로 불가능하다.

시간복잡도는 리스트를 만드는 시간에 의존하므로 $O(n)$이다.

참고

실행속도와 메모리는 다음과 같다.

정확성 테스트 결과

효율성 테스트 결과

 

산술-기하 평균 부등식 영어 위키

https://en.wikipedia.org/wiki/AM%E2%80%93GM_inequality

AM–GM inequality - Wikipedia