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

파이썬으로 숫자 집합의 부분 합 제곱 최댓값 구하기 – 각도 정렬과 그리디 탐색

문제 이해하기

nums1과 nums2라는 두 개의 배열이 주어지고, 두 배열은 모두 같은 개수의 원소 N개를 가지고 있다고 가정해 보겠습니다. 이제 1부터 N까지의 자연수로 이루어진 집합 S를 생각해 봅니다. 우리가 구해야 하는 값은 다음 식의 최댓값입니다.

(nums1[i1] + nums1[i2] + … + nums1[ik])² + (nums2[i1] + nums2[i2] + … + nums2[ik])²

여기서 {i1, i2, …, ik}는 집합 S의 공집합이 아닌 임의의 부분 집합입니다. 즉, 인덱스를 어떻게 선택하느냐에 따라 두 배열에서 선택된 값들의 합을 각각 제곱해 더한 값 중 가장 큰 것을 찾는 문제입니다.

예시

입력이 다음과 같다고 해보겠습니다.

nums1 = [-1, 6], nums2 = [5, 4]

이 경우 출력은 106이 됩니다. 가능한 조합을 살펴보면 다음과 같습니다.

  • (-1)² + (5)² = 26
  • (6)² + (4)² = 50
  • (-1 + 6)² + (5 + 4)² = 106

해결 접근 방법

이 문제는 각 인덱스 i를 좌표평면상의 한 점, 즉 벡터 (nums1[i], nums2[i])로 생각하면 기하학적으로 접근할 수 있습니다. 부분 집합의 선택 결과는 곧 선택된 벡터들의 합벡터가 되고, 우리가 최대화하려는 값은 바로 그 합벡터 길이의 제곱입니다.

핵심 아이디어는 다음과 같습니다.

  • 각 쌍 (nums1[i], nums2[i])를 원점 기준 각도(atan2) 순으로 정렬합니다.
  • 각 벡터를 시작점으로 삼아, 정렬된 순서를 순환 구조로 따라가며 벡터를 하나씩 더해갑니다.
  • 더했을 때 제곱합이 커지는 동안만 계속 누적하고, 더 이상 커지지 않으면 멈춥니다. 이 과정을 정방향과 역방향 두 방향으로 모두 수행합니다.
  • 탐색 과정에서 얻은 최댓값을 반환합니다.

이 그리디 전략이 성립하는 이유는, 최적의 부분 집합의 합벡터가 항상 각도 순으로 정렬된 연속 구간에 대응되기 때문입니다. 덕분에 모든 부분 집합을 확인하는 완전 탐색(O(2N)) 대신 O(N²) 시간 안에 답을 구할 수 있습니다.

알고리즘을 단계별로 정리하면 다음과 같습니다.

  1. vs := 0부터 nums1의 크기 - 1까지 각 i에 대해 (nums1[i], nums2[i]) 쌍의 리스트를 만듭니다.
  2. vs를 각 원소 v에 대해 atan2(v[1], v[0]) 값을 기준으로 정렬합니다.
  3. best := 0으로 초기화합니다.
  4. i를 0부터 vs의 크기 - 1까지 반복하면서:
    • u := vs[i], l := u[0]² + u[1]²로 초기화합니다.
    • (vs + vs)[i+1 : i+len(vs)] 구간의 각 v에 대해 t1 = u + v, t2 = t1의 제곱합을 계산하고, t2 ≥ l이면 u와 l을 갱신합니다.
    • l > best이면 best를 갱신합니다.
    • 같은 과정을 역순(reversed)으로 한 번 더 수행합니다.
  5. best를 반환합니다.

구현 예제

다음 파이썬 코드를 통해 더 잘 이해할 수 있습니다.

from math import atan2

def solve(nums1, nums2):
    vs = zip(nums1, nums2)
    vs = sorted(vs, key=lambda v: atan2(v[1], v[0]))

    best = 0
    for i in range(len(vs)):
        # 정방향 탐색
        u = vs[i]
        l = u[0]*u[0] + u[1]*u[1]
        for v in (vs + vs)[i+1:(i + len(vs))]:
            t1 = (u[0]+v[0], u[1]+v[1])
            t2 = t1[0]*t1[0] + t1[1]*t1[1]
            if t2 >= l:
                u = t1
                l = t2
        if l > best:
            best = l

        # 역방향 탐색
        u = vs[i]
        l = u[0]*u[0] + u[1]*u[1]
        for v in reversed((vs + vs)[i+1:(i + len(vs))]):
            t1 = (u[0]+v[0], u[1]+v[1])
            t2 = t1[0]*t1[0] + t1[1]*t1[1]
            if t2 >= l:
                u = t1
                l = t2
            if l > best:
                best = l
    return best

nums1 = [-1, 6]
nums2 = [5, -4]
print(solve(nums1, nums2))

실행 결과

입력:

[-1, 6], [5, -4]

출력:

52

위 코드에서 nums2 = [5, -4]를 사용하면 세 가지 가능한 부분 집합의 결과는 각각 26((-1)² + 5²), 52(6² + (-4)²), 26((-1+6)² + (5-4)²)이므로, 최댓값인 52가 출력됩니다. 만약 nums2 = [5, 4]였다면 앞서 본 것처럼 106이 최댓값이 됩니다.

마무리

이처럼 벡터를 각도 순으로 정렬한 뒤 그리디하게 누적하는 방식을 활용하면, 모든 부분 집합을 일일이 확인하지 않고도 O(N²) 만에 표현식의 최댓값을 효율적으로 찾을 수 있습니다. 좌표 기하학의 관점을 알고리즘 문제에 적용하는 좋은 예시이니, 비슷한 유형의 부분 집합 최적화 문제를 만났을 때 꼭 떠올려 보시기 바랍니다.