Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 크기, 총합, 최댓값 조건을 만족하는 중복 없는 배열 구성하기

문제 개요

크기를 나타내는 변수 N, 배열 원소들의 총합을 나타내는 변수 SUM, 그리고 "배열의 어떤 원소도 K보다 클 수 없다"는 조건이 주어진다고 가정해 봅시다. 이때 우리가 찾아야 하는 것은 모든 원소가 서로 다른(중복이 없는) 배열입니다. 만약 주어진 조건을 만족하는 배열이 존재하지 않는다면 -1을 반환해야 합니다.

예를 들어 입력이 N = 4, SUM = 16, K = 9라면 출력은 [1, 2, 4, 9]가 됩니다.

접근 방법

이 문제의 핵심은 서로 다른 원소로 만들 수 있는 합계의 범위를 먼저 파악하는 것입니다.

  • 최소 합(minimum_sum): 1부터 N까지의 합, 즉 (N × (N + 1)) / 2 입니다. 중복 없는 양의 정수로 만들 수 있는 가장 작은 합계입니다.
  • 최대 합(maximum_sum): K, K-1, ..., K-N+1처럼 K 이하의 서로 다른 정수 N개로 만들 수 있는 합, 즉 (N × K) − (N × (N − 1)) / 2 입니다.

SUM이 이 범위를 벗어나면 조건을 만족하는 배열이 존재하지 않으므로 -1을 반환합니다. 범위 안에 있다면 다음 절차로 답을 구성합니다.

  • res를 [0, 1, 2, ..., N]으로 초기화합니다. (인덱스 0은 자리 표시용이며, 실제 배열은 인덱스 1~N에 해당합니다.)
  • 현재 합계를 최소 합으로 설정하고, i를 N부터 1까지 감소시키며 탐색합니다.
  • x = 현재 합 + (K − i) 를 계산했을 때 x가 SUM보다 작으면, res[i]를 K로 교체하고 합계를 늘린 뒤 K를 1 감소시킵니다.
  • x가 SUM 이상이면, res[i]에 남은 차이(SUM − 현재 합)를 더해 목표 합계를 정확히 맞춘 뒤 반복을 종료합니다.

구현 예제

아래 구현을 통해 동작 방식을 더 잘 이해할 수 있습니다. 원본 코드에서 정수 나눗셈(/) 때문에 결과에 소수점이 붙던 부분은 정수 나눗셈(//)으로, 내장 함수 sum과 이름이 겹치던 변수는 total로 바꿔 가독성을 개선했습니다.

def get_arr(N, SUM, K):
    minimum_sum = (N * (N + 1)) // 2
    maximum_sum = (N * K) - (N * (N - 1)) // 2
    if minimum_sum > SUM or maximum_sum < SUM:
        return -1

    res = [i for i in range(N + 1)]
    total = minimum_sum
    i = N

    while i >= 1:
        x = total + (K - i)
        if x < SUM:
            total += (K - i)
            res[i] = K
            K -= 1
        else:
            res[i] += (SUM - total)
            total = SUM
            break
        i -= 1

    return res

N = 4
SUM = 16
K = 9
print(get_arr(N, SUM, K))

입력

N = 4, SUM = 16, K = 9

출력

[0, 1, 2, 4, 9]

인덱스 0은 자리 표시용이므로, 실제 정답 배열은 인덱스 1부터 N까지의 값인 [1, 2, 4, 9]입니다.

동작 과정 살펴보기

  • 초기 상태: res = [0, 1, 2, 3, 4], 현재 합계 = 10 (최소 합)
  • i = 4일 때: x = 10 + (9 − 4) = 15 < 16 이므로 res[4]를 9로 교체, 합계는 15, K는 8이 됩니다.
  • i = 3일 때: x = 15 + (8 − 3) = 20 ≥ 16 이므로 res[3]에 남은 차이 16 − 15 = 1을 더해 4로 만들고, 합계가 16에 도달했으므로 종료합니다.

최종적으로 [0, 1, 2, 4, 9]가 반환되며, 모든 원소는 서로 다르고 K = 9 이하이면서 총합이 정확히 16입니다. 이 알고리즘은 배열을 한 번만 순회하므로 시간 복잡도는 O(N)으로 매우 효율적입니다.