파이썬으로 행 합과 열 합이 주어진 유효한 행렬 찾기

문제 개요

음수가 아닌 값으로 이루어진 두 배열 rowSum과 colSum이 주어집니다. 여기서 rowSum[i]는 2차원 행렬의 i번째 행에 있는 원소들의 합을, colSum[j]는 j번째 열에 있는 원소들의 합을 의미합니다. 목표는 이 두 조건을 모두 만족하는, 크기가 (rowSum의 길이 × colSum의 길이)인 음수가 아닌 값의 행렬을 하나 찾는 것입니다.

예를 들어 입력이 rowSum = [13, 14, 12], colSum = [9, 13, 17]이라면 다음과 같은 행렬이 정답이 될 수 있습니다.

실제로 각 행의 합은 9+4+0=13, 0+9+5=14, 0+0+12=12이고, 각 열의 합은 9+0+0=9, 4+9+0=13, 0+5+12=17로 조건을 정확히 만족합니다.

접근 방법: 그리디(Greedy) 알고리즘

이 문제는 그리디 방식으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

• 매 단계마다 아직 처리되지 않은 행 합과 열 합 중 더 작은 값을 선택합니다.
• 선택된 값을 해당 행과 열이 교차하는 위치의 행렬 원소로 배정합니다.
• 그러면 행 또는 열 중 하나의 합이 완전히 소진되므로, 해당 행(또는 열)은 더 이상 고려하지 않습니다.
• 모든 행과 열의 합이 처리될 때까지 이 과정을 반복합니다.

구체적인 알고리즘 단계는 다음과 같습니다.

• 빈 행렬(matrix)을 생성하고, 방문 기록용 집합(visited)을 준비합니다.
• minimum() 함수를 정의합니다. 이 함수는 남은 행 합 배열 r과 열 합 배열 c를 인자로 받습니다.
• min_total을 무한대로, type을 빈 문자열로 초기화합니다.
• r의 모든 원소를 순회하며 r[i]가 min_total보다 작으면 index, type('row'), min_total을 갱신합니다.
• c의 모든 원소를 순회하며 c[i]가 min_total보다 작으면 같은 방식으로 갱신합니다.
• type이 'row'라면 r[index]를 무한대로 설정한 뒤, 무한대가 아니면서 min_total 이상인 첫 번째 열 원소 c[i]를 찾아 c[i]에서 min_total을 빼고 matrix[index][i]에 min_total을 저장한 후 루프를 종료합니다.
• type이 'col'이라면 위와 대칭적으로 c[index]를 무한대로 설정하고 행 방향으로 동일하게 처리합니다.
• (index, type) 쌍을 visited에 추가합니다.
• 메인 루프에서는 visited의 크기가 len(r)+len(c)와 같아질 때까지 minimum(r, c)를 호출합니다.
• 완성된 matrix를 반환합니다.

파이썬 구현 코드

다음 구현을 통해 동작을 더 잘 이해할 수 있습니다.

def solve(r, c):
    matrix = [[0]*len(c) for _ in range(len(r))]
    visited = set()

    def minimum(r, c):
        min_total = float('inf')
        type = ''
        for i in range(len(r)):
            if(r[i] < min_total):
                index = i
                type = 'row'
                min_total = r[i]

        for i in range(len(c)):
            if(c[i] < min_total):
                min_total = c[i]
                type = 'col'
                index = i

        if(type == 'row'):
            r[index] = float('inf')
            for i in range(len(c)):
                if(c[i] != float('inf') and c[i] >= min_total):
                    c[i] -= min_total
                    matrix[index][i] = min_total
                    break

        if(type == 'col'):
            c[index] = float('inf')
            for i in range(len(r)):
                if(r[i] != float('inf') and r[i] >= min_total):
                    r[i] -= min_total
                    matrix[i][index] = min_total
                    break

        visited.add((index, type))

    while len(visited) != len(r) + len(c):
        minimum(r, c)

    return matrix

rowSum = [13, 14, 12]
colSum = [9, 13, 17]
print(solve(rowSum, colSum))

입력

[13, 14, 12], [9, 13, 17]

출력

[[9, 4, 0], [0, 9, 5], [0, 0, 12]]

복잡도 분석

행의 개수를 m, 열의 개수를 n이라 하면, while 루프는 최대 m+n번 실행되고 각 호출에서 배열 전체를 순회하므로 시간 복잡도는 O((m+n)²)입니다. 결과 행렬 저장에는 O(m×n)의 공간이 필요합니다. 참고로, 이미 소진된 행·열을 리스트에서 제거하는 방식으로 최적화하면 O(m×n)까지 개선할 수 있습니다.