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

파이썬으로 n번의 연산 후 얻을 수 있는 최대 점수 구하기

문제 설명

크기가 2×n인 배열 nums가 주어집니다. 우리는 이 배열에 대해 정확히 n번의 연산을 수행해야 하며, i번째 연산(1부터 시작하는 인덱스)에서는 다음과 같은 작업을 진행합니다.

  • 배열에서 두 개의 원소 x와 y를 선택합니다.
  • i × gcd(x, y) 만큼의 점수를 획득합니다. 여기서 gcd는 최대공약수를 의미합니다.
  • 선택한 두 원소 x와 y를 배열에서 제거합니다.

n번의 연산을 모두 수행한 후 얻을 수 있는 최대 점수를 구하는 것이 목표입니다.

예를 들어, 입력이 nums = [6,2,1,5,4,3]이라면 출력은 14가 됩니다. 최적의 선택 순서는 다음과 같습니다.

(1 × gcd(1, 5)) + (2 × gcd(2, 4)) + (3 × gcd(3, 6)) = 1 + 4 + 9 = 14

해결 방법: 비트마스크 DP + 메모이제이션

이 문제는 어떤 원소들이 이미 사용되었는지 상태를 추적해야 하므로, 비트마스크(bitmask) 기반의 동적 계획법(DP)과 메모이제이션을 활용하면 효율적으로 해결할 수 있습니다.

핵심 아이디어

  • n := nums 배열의 크기로 설정합니다.
  • dp := 크기가 2^n인 배열을 생성하고 모든 값을 -1로 초기화합니다. 각 인덱스는 '현재까지 선택된 원소들의 조합'을 나타내는 비트마스크입니다.
  • dfs(mask, t) 함수를 정의합니다. mask는 이미 사용된 원소들을 나타내고, t는 현재 연산 번호입니다.

dfs() 함수의 동작 과정

  1. 종료 조건: mask가 (2^n - 1)과 같다면, 즉 모든 원소를 사용했다면 0을 반환합니다.
  2. 메모이제이션 확인: dp[mask]가 -1이 아니라면, 이미 계산된 값이므로 그 값을 그대로 반환합니다.
  3. ma := 0으로 초기화합니다.
  4. i를 0부터 n-1까지 반복하며, mask에서 i번째 비트가 이미 켜져 있으면(원소가 사용되었다면) 건너뜁니다.
  5. j를 i+1부터 n-1까지 반복하며, 마찬가지로 j번째 비트가 켜져 있으면 건너뜁니다.
  6. next := dfs(mask에 i번째와 j번째 비트를 추가한 값, t+1) + gcd(nums[i], nums[j]) × t 를 계산합니다.
  7. ma := next와 ma 중 더 큰 값으로 갱신합니다.

모든 탐색이 끝나면 dp[mask] := ma를 저장하고 반환합니다. 메인 호출부에서는 dfs(0, 1)의 결과를 반환하면 됩니다.

구현 예시

아래 코드를 통해 실제 동작을 더 잘 이해할 수 있습니다.

from math import gcd

def solve(nums):
    n = len(nums)
    dp = [-1] * (1 << n)

    def dfs(mask, t):
        # 모든 원소를 사용한 경우 종료
        if mask == (1 << n) - 1:
            return 0
        # 이미 계산된 상태라면 저장된 값 반환
        if dp[mask] != -1:
            return dp[mask]
        ma = 0
        for i in range(n):
            # i번째 원소가 이미 사용되었다면 건너뛰기
            if (1 << i) & mask:
                continue
            for j in range(i + 1, n):
                # j번째 원소가 이미 사용되었다면 건너뛰기
                if (1 << j) & mask:
                    continue
                next = dfs(mask | (1 << i) | (1 << j), t + 1) + gcd(nums[i], nums[j]) * t
                ma = max(next, ma)
        dp[mask] = ma
        return dp[mask]

    return dfs(0, 1)

nums = [6, 2, 1, 5, 4, 3]
print(solve(nums))

입력

[6,2,1,5,4,3]

출력

14

복잡도 분석

  • 시간 복잡도: O(2^n × n²) — 가능한 모든 비트마스크 상태에 대해 두 원소를 선택하는 조합을 탐색하지만, 메모이제이션 덕분에 같은 상태는 한 번만 계산됩니다.
  • 공간 복잡도: O(2^n) — 비트마스크 상태별 결과를 저장하는 dp 배열이 필요합니다.

이처럼 비트마스크 DP를 활용하면 원소의 조합을 효율적으로 관리하면서 최적의 점수를 보장할 수 있습니다.