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

파이썬으로 합이 k 이하인 직사각형의 최대 합 구하기

문제 개요

2차원 행렬과 하나의 값 k가 주어졌을 때, 합이 k 이하(sum ≤ k)가 되는 직사각형 영역 중에서 가장 큰 합을 찾는 것이 목표입니다. 여기서 직사각형은 행렬 안에서 연속된 행과 열로 이루어진 부분 행렬을 의미합니다.

예를 들어 다음과 같은 행렬이 있다고 가정해 보겠습니다.

5-2
710

k = 15일 때 정답은 12입니다. 첫 번째 열로 이루어진 직사각형 [5, 7]의 합이 12로, 15를 넘지 않으면서 만들 수 있는 최대 합이기 때문입니다.

풀이 접근 방법

이 문제는 다음과 같은 아이디어로 해결할 수 있습니다. 먼저 위쪽 행(i1)과 아래쪽 행(i2)을 고정하여 두 행 사이의 영역을 하나의 1차원 배열로 압축합니다. 그런 다음 누적합(prefix sum)을 활용하여, 현재까지의 합에서 k를 뺀 값보다 큰 이전 누적합 중 가장 작은 것을 찾으면, 그 차이가 곧 k 이하인 직사각형의 합이 됩니다. 이 과정을 가능한 모든 행 조합에 대해 반복하면 전체 최댓값을 구할 수 있습니다.

단계별 알고리즘

  • n := 행렬 a의 행 개수
  • m := 행렬 a의 열 개수
  • ans := 충분히 작은 값으로 초기화
  • i1을 0부터 n-1까지 반복합니다.
    • row := 크기가 m이고 0으로 채워진 리스트 생성
    • i2를 i1부터 n-1까지 반복합니다.
      • j를 0부터 m-1까지 반복하며 row[j]에 a[i2][j]를 더합니다.
      • s := 새로운 집합(set)을 만들고 0을 삽입합니다.
      • sum := 0
      • j를 0부터 m-1까지 반복합니다.
        • sum := sum + row[j]
        • temp := s에 있는 원소 중 (sum − k)보다 큰 값들의 목록
        • temp가 비어 있지 않으면:
          • u := temp의 최솟값
          • ans := ans와 (sum − u) 중 더 큰 값으로 갱신
        • s에 sum을 삽입합니다.
  • ans를 반환합니다.

여기서 (sum − u)는 현재 위치까지의 부분 배열 합 중 k 이하가 되는 최대값을 의미합니다. 이전 누적합 u가 (sum − k)보다 크면 그 차이가 반드시 k 이하가 되므로 문제의 조건을 만족하게 됩니다.

구현 예시

class Solution:
    def solve(self, a, k):
        n = len(a)
        if n == 0:
            return 0
        m = len(a[0])
        ans = -999999
        for i1 in range(n):
            row = [0]*m
            for i2 in range(i1, n):
                for j in range(m):
                    row[j] += a[i2][j]
                s = set()
                s.add(0)
                sum = 0
                for j in range(m):
                    sum += row[j]
                    temp = [e for e in s if e > (sum - k)]
                    if len(temp) > 0:
                        u = min(temp)
                        ans = max(ans, sum - u)
                    s.add(sum)
        return ans

ob = Solution()
matrix = [
    [5, -2],
    [7, 10]
]
k = 15
print(ob.solve(matrix, k))

입력

[
[5, -2],
[7, 10]
], 15

출력

12

시간 복잡도

행 쌍을 고정하는 데 O(n²)이 걸리고, 각 단계에서 집합을 선형 탐색하므로 전체 시간 복잡도는 최악의 경우 O(n² × m²)입니다. 집합 대신 정렬된 리스트와 이진 탐색(bisect 모듈)을 함께 사용하면 O(n² × m log m)으로 성능을 개선할 수 있습니다.