[백준 17123번] 파이썬 - 배열 놀이

백준 17123 - 배열 놀이

시간 제한 2초, 메모리 제한 512MB

# 조건

• N개의 행과 N개의 열로 구성된 2차원 정수 배열 A가 있다.
• A[r, c]는 r번째 행 c번째 열에 위치한 원소의 값을 나타낸다.
• 이 배열에 총 M번의 연산을 적용하는 배열 놀이를 생각해보자.
• 각 연산에 대해 1 <= r1 <= r2 <= N, 1 <= c1 <= c2 <= N, -1,000 <= v <= 1,000을 만족하는 다섯 개의 정수 파라미터 (r1, c1, r2, c2, v)가 주어지는데, 이 연산은 (r1, c1)부터 (r2, c2)까지 사각형 영역에 속한 A[r, c]의 값을 v만큼 더하는 것이다.
• 예를 들어 N = 3 그리고 A = [ [ 1 2 3 ] [ 4 5 6 ] [ 7 8 9 ] ] 인 2차원 배열을 생각해보자. 아래와 같은 3개의 연산을 순서대로 적용한다고 하자.
  • 연산 1: (r1 = 1, c1 = 1, r2 = 2, c2 = 3, v = 3)
  • 연산 2: (r1 = 2, c1 = 2, r2 = 3, c2 = 2, v = -5)
  • 연산 3: (r1 = 1, c1 = 1, r2 = 3, c2 = 2, v = 1)
• 연산을 적용하기 전 A는 아래와 같다.

A = [ [ 1 2 3 ]
      [ 4 5 6 ]
      [ 7 8 9 ] ].

• 연산 1을 적용하면 첫 두 행에 포함된 여섯 개의 원소 값이 바뀌게 되어 아래와 같아진다 (굵게 표시된 수들의 값이 바뀌었다):

A = [ [ **4 5 6** ]
      [ **7 8 9** ]
      [ 7 8 9 ] ].

• 연산 2를 적용한 후:

A = [ [ 4 5 6 ]
      [ 7 **3** 9 ]
      [ 7 **3** 9 ] ].

• 연산 3을 적용한 후:

A = [ [ **5 6** 6 ]
      [ **8 4** 9 ]
      [ **8 4** 9 ] ].

• 이렇게 총 3개의 연산을 모두 적용하고 난 후, 당신은 마지막에 얻은 배열의 각 행의 원소들의 합과 각 열의 원소들의 합을 구하고 싶다.
• 위의 예제의 경우, 행의 원소들의 합은 [17, 21, 21] 이며 (1번 행 부터 3번 행까지) 열의 원소들의 합은 [21, 14, 24] 이다 (1번 열 부터 3번 열 까지).
• 입력으로 N, M, 2차원 배열 A, 그리고 M개의 연산이 주어졌을 때, 배열에 연산을 모두 적용한 후 해당 배열의 행의 원소들의 합과 열의 원소들의 합을 구하는 프로그램을 작성하시오.

입력

• 첫 줄에 테스트 케이스의 수 T가 주어진다 (1 <= T <= 10).
• 각 테스트 케이스에 대하여: 첫 줄에 두 개의 정수 N과 M이 공백으로 구분되어 주어진다 (1 <= N <= 1,000 과 1 <= M <= 1,000 을 만족한다).
• 다음 N줄에 걸쳐서 2차원 배열 A가 주어지는데, i-번째 줄이 i-번째 행을 나타낸다. 각 줄의 j번째 정수는 j-번째 열의 원소 값을 나타낸다.
  • 배열 A의 각 원소의 값은 1이상 1,000 이하의 정수이다.
• 다음 M줄에 걸쳐서 각 줄에 다섯 개의 정수 r1, c1, r2, c2, v가 공백으로 구분되어 주어진다.
• 항상 1 <= r1 <= r2 <= N 그리고 1 <= c1 <= c2 <= N 그리고 -1,000 <= v <= 1,000 을 만족한다.

출력

• 각 테스트 케이스에 대해 두 줄에 걸쳐서 정답을 출력해야 한다.
• 두 줄 중 첫 줄에는 N개의 정수로 표현된 각 행의 합을 공백으로 구분하여 출력하고 (1번 행 부터 N번 행 까지)
• 둘째 줄에는 N개의 정수로 표현된 각 열의 합을 공백으로 구분하여 출력한다 (1번 열 부터 N번 열 까지).

# 접근 방법

• 모든 경우를 실시간으로 더하고 빼준다면 한 테스트 케이스 당 최대 1000 * 1000 * 1000의 시간 복잡도가 걸린다.
• 따라서 누적합을 이용해줄 건데, 시작 행부터 도착 행까지 모든 시작 열에 더하거나 뺄 값을 기록하고, 도착 열 +1에 -val을 기록하여 구할 수 있다.
  • 우선 주어진 배열 크기에 행과 열을 한 줄씩 더 추가한 크기로 prefix 배열을 만들어준다.
  • ex) 1,1 에서 3,2 까지 3씩 더한다면
  • (1, 1), (2, 1), (3, 1)에 3을 기록
  • 도착 열 +1 인 3에 -3 기록 => (1, 3), (2, 3), (3, 3)
  • 왼쪽에서 부터 누적합을 구하면 3 3 3+(-3) 과 같이 되므로 목표했던 2열까지만 더해진다.
• 다만 이 경우에도 모든 행에 기록을 해주므로 2 * 1000 * 1000의 시간 복잡도가 나올 수 있기 때문에 위의 방식을 조금 변경해주었다.
• 모든 행에 기록하는 것이 아닌 열과 마찬가지로 끝 행에만 기록을 해준다.
• 즉 (si, sj), (ni, nj) 가 시작 구간과 끝 구간으로 주어졌을 때
  • (si, sj) => 시작행, 시작열 에는 val 값을
  • (si, nj + 1) => 시작행, 끝 열 + 1 에는 -val 값을
  • (ni +1, sj) => 끝 행 + 1, 시작 열에는 -val 값을
  • (ni +1, nj+1) => 끝 행 + 1, 끝 열 +1에는 val 값을 기록해준다.
• 예를 들어, 1, 1에서 3, 2까지 3씩 더해야된다면
  • 1, 1 에는 3
  • 1, 3에는 -3
  • 4, 1에는 -3
  • 4, 3에는 3을 기록하면 prefix의 값은 아래와 같이 되고
  • 모든 행을 왼쪽에서 오른쪽으로 누적 합을 구해준 뒤 모든 열을 위에서 아래로 누적 합을 한 번더 구해준다면 처음 방식과 같은 결과를 얻을 수 있다.

# 시작 구간과 끝 구간의 행, 열에 표시
[
 [3, 0, -3, 0]
 [0, 0, 0, 0]
 [0, 0, 0, 0]
 [-3, 0, 3, 0]
]
# 가로로 더한 뒤
[
 [3, 3, 0, 0]
 [0, 0, 0, 0]
 [0, 0, 0, 0]
 [-3, -3, 0, 0]
]
# 세로로 더한 뒤
# 처음 원하는 구간에 원하는 value만큼 더해진걸 볼 수 있음
[
 [3, 3, 0, 0]
 [3, 3, 0, 0]
 [3, 3, 0, 0]
 [0, 0, 0, 0]
]

• 이후 원본 배열에 구해 놓은 누적 합을 더해주면 되는데 우리가 원하는 출력은 각 행의 합과, 각 열의 합을 출력하는 것이다.
• 따라서 배열에 값을 더하는 동시에 row_temp 에 행의 누적 합을 더해주고 col_temp[j]에 각 열의 합을 더해준다.

import sys
input = sys.stdin.readline

# 누적합 구해주기
# (si, sj), (ni, nj) 가 시작 구간과 끝 구간으로 주어졌을 때
#     - (si, sj) => 시작행, 시작열 에는 val 값을
#     - (si, nj + 1) => 시작행, 끝 열 + 1 에는 -val 값을
#     - (ni +1, sj) => 끝 행 + 1, 시작 열에는 -val 값을
#     - (ni +1, nj+1) => 끝 행 + 1, 끝 열 +1에는 val 값을 기록해준다.
# 시작 구간과 끝 구간의 행, 열에 표시
# [
#  [3, 0, -3, 0]
#  [0, 0, 0, 0]
#  [0, 0, 0, 0]
#  [-3, 0, 3, 0]
# ]
# # 가로로 더한 뒤
# [
#     #  [3, 3, 0, 0]
#     #  [0, 0, 0, 0]
#     #  [0, 0, 0, 0]
#     #  [-3, -3, 0, 0]
# ]
# # 세로로 더한 뒤
# [
#  [3, 3, 0, 0]
#  [3, 3, 0, 0]
#  [3, 3, 0, 0]
#  [0, 0, 0, 0]
# ]
# # 처음 원하는 구간에 원하는 value만큼 더해진걸 볼 수 있음

def solve():
    for _ in range(int(input())):
        N, M = map(int, input().split())
        arr = [[*map(int, input().split())] for _ in range(N)]
        prefix = [[0] * (N+1) for _ in range(N+1)]
        # 구간의 시작과 끝에 기록
        for _ in range(M):
            si, sj, ni, nj, val = map(lambda x:int(x)-1, input().split())
            val += 1
            prefix[si][sj] += val
            prefix[si][nj+1] += -val
            prefix[ni+1][sj] += -val
            prefix[ni+1][nj+1] += val

        # 누적합 행 기준 한번 열기준 한번
        for i in range(N+1):
            for j in range(1, N+1):
                prefix[i][j] += prefix[i][j-1]
        for j in range(N+1):
            for i in range(1, N+1):
                prefix[i][j] += prefix[i-1][j]

        # 각 행과 열의 합을 출력해야 되므로 누적합을 원본 배열에 더해주면서 저장해주기
        col_temp = [0] * N
        for i in range(N):
            row_temp = 0
            for j in range(N):
                arr[i][j] += prefix[i][j]
                row_temp += arr[i][j]
                col_temp[j] += arr[i][j]
            print(row_temp, end=' ')
        print()
        print(*col_temp)
solve()

• 처음에 잘 살펴보지 않고, 원본 배열에 모든 연산을 수행한 값을 출력해야 되는 걸로 착각하였다;
• 불필요한 누적 합을 구하는 부분이 들어갔지만 모든 연산에 대해 누적 합을 효율적으로 구하였기에 통과할 수 있었다.
• 다른 분들의 풀이를 보면 미리 원본 배열의 행과 열의 합을 리스트로 구해놓고 주어진 구간의 행, 열의 개수 * value 값을 더해가는 식으로 구해주었다.

T = int(input())  


for _ in range(T):  
    N, M = map(int, input().split())  
    arr = [list(map(int, input().split())) for _ in range(N)]  
    rowsum = [sum(a) for a in arr]  
    colsum = [sum(a) for a in zip(*arr)]  

    for _ in range(M):  
        r1, c1, r2, c2, v = map(int, input().split())  
        r = r2 - r1 + 1  # 2  
        c = c2 - c1 + 1  # 3  
        for i in range(r1 - 1, r2):  
            rowsum[i] += c * v  
        for i in range(c1 - 1, c2):  
            colsum[i] += r * v  
    print(' '.join(map(str, rowsum)))  
    print(' '.join(map(str, colsum)))