[ 백준 17144번 미세먼지 안녕! ] Python 코드 (Simulation)

1. 문제

https://www.acmicpc.net/problem/17144

17144번: 미세먼지 안녕!

---

2. 풀이

문제가 어려운 것은 아니었지만 Python3로 돌리면 시간초과가 발생했다.

PyPy3로는 통과되었지만 뭔가 만족스럽지 않아서 Python3로 통과될 때까지 코드를 최적화했다.

 

그래서 어떻게 겨우겨우 Python3로도 통과를 했다!!

상위코드를 몇 개를 살펴 보았는데, 내가 제출한 코드와 비슷한데도 2배 정도 시간차이가 났다. 어째서인지 잘 모르겠다.

 

내가 적용한 최적화 방법들은 아래와 같다.

2.1. 짧은 for문 풀어헤치기

보통 BFS나 DFS에서 2차원 배열을 탐색할 때, 보통 세번째 for문처럼 4방향을 탐색한다.

그런데 교수님께 들은 바로는, 코드가 for문에 진입하는데 내부 로직상 추가시간이 발생한다고 한다.

그래서 정말로 코드를 최적화해야 하는 경우에는 짧은 for문은 풀어버리라고 말씀하셨다. 

# 확산
arr_tmp = [[0] * C for _ in range(R)]
for i in range(R):
    for j in range(C):
        if arr[i][j] > 4:	# 문제조건상 4이하면 실질적으로 확산이 안됨
            cnt = 0
            for di, dj in ((1, 0), (0, 1), (-1, 0), (0, -1)):
                ni = i + di
                nj = j + dj
                if 0 <= ni < R and 0 <= nj < C and arr[ni][nj] != -1:
                    arr_tmp[ni][nj] += arr[i][j] // 5
                    cnt += 1
            arr[i][j] = arr[i][j] - (arr[i][j] // 5) * cnt

위 코드의 for문을 아래처럼 풀었고, if문은 문제 조건에 적합하게 수정하였다.

오른쪽(right)방향 확산 시 if문이 다른 방향들의 if문과 조금 다르다.

왜냐하면 공기청정기가 0열에 있으므로, 오른쪽 방향 확산 시에는 공기청정기의 위치를 고려할 필요 없기 때문이다.

# 확산
arr_tmp = [[0] * C for _ in range(R)]
for i in range(R):
    for j in range(C):
        if arr[i][j] > 4:	# 문제조건상 4이하면 실질적으로 확산이 안됨
            cnt = 0
            if i < R - 1 and arr[i+1][j] != -1:     # down
                arr_tmp[i + 1][j] += arr[i][j] // 5
                cnt += 1
            if i > 0 and arr[i-1][j] != -1:         # up
                arr_tmp[i - 1][j] += arr[i][j] // 5
                cnt += 1
            if j < C - 1:                           # right
                arr_tmp[i][j + 1] += arr[i][j] // 5
                cnt += 1
            if j > 0 and arr[i][j-1] != -1:         # left
                arr_tmp[i][j - 1] += arr[i][j] // 5
                cnt += 1
            arr[i][j] -= arr[i][j] // 5 * cnt

2.2. 동일 연산 반복하지 않기

또한 바로 위의 코드를 보면, arr[i][j] // 5가 반복적으로 쓰이고 있다.

이런 경우 처음의 연산값을 저장해서 쓰면, 불필요한 연산을 줄일 수 있다.

추가적으로 마지막 arr[i][j] -= arr[i][j] // 5 * cnt 를 없애고, 대신 매 조건문마다 arr[i][j] -= spread를 넣어주었다.

시간복잡도상의 이점은 미미하나, 코드가 전반적으로 보기 좋아졌다고 생각한다.

# 확산
arr_tmp = [[0] * C for _ in range(R)]
for i in range(R):
    for j in range(C):
        if arr[i][j] > 4:                       
            spread = arr[i][j] // 5	            # 첫번째 연산 시의 값을 저장한다.
            if i < R - 1 and arr[i+1][j] != -1:     # down
                arr_tmp[i + 1][j] += spread
                arr[i][j] -= spread
            if i > 0 and arr[i-1][j] != -1:         # up
                arr_tmp[i - 1][j] += spread
                arr[i][j] -= spread
            if j < C - 1:                           # right
                arr_tmp[i][j + 1] += spread
                arr[i][j] -= spread
            if j > 0 and arr[i][j-1] != -1:         # left
                arr_tmp[i][j - 1] += spread
                arr[i][j] -= spread
        arr_tmp[i][j] += arr[i][j]

 

2.3. 조건 연산 횟수 줄이기

위의 if 조건문들은 대부분이 and 연산을 가지고 있다.

and 연산은 전자 조건이 True일 경우 단축평가를 지원하지 않는다.

그리고 전자의 조건을 보면 대부분의 경우 True일 것임을 짐작할 수 있다.

따라서 위의 if문들은 대부분 조건문 판단을 2회씩 하고 있다.

 

그런데 보면 and의 후자 조건들은 모두 공기청정기의 위치 때문에 발생한다.

공기청정기는 0열에만 존재하기 때문에 2열 이상의 위치에서는 이 조건문들이 불필요하다.

따라서 0~1열과 2열 이상을 분리한다면, 코드는 길어지겠지만 조건 판단 횟수를 줄이기 때문에 실행속도는 빨라질 것이다.

# 확산
arr_tmp = [[0] * C for _ in range(R)]
for i in range(R):
    for j in range(2):                      # 0, 1 열
        if arr[i][j] > 4:                       # 문제조건상 4이하면 실질적으로 확산이 안됨
            spread = arr[i][j] // 5
            if i < R - 1 and arr[i+1][j] != -1:     # down
                arr_tmp[i + 1][j] += spread
                arr[i][j] -= spread
            if i > 0 and arr[i-1][j] != -1:         # up
                arr_tmp[i - 1][j] += spread
                arr[i][j] -= spread
            if j < C - 1:                           # right
                arr_tmp[i][j + 1] += spread
                arr[i][j] -= spread
            if j > 0 and arr[i][j-1] != -1:         # left
                arr_tmp[i][j - 1] += spread
                arr[i][j] -= spread
        arr_tmp[i][j] += arr[i][j]
    for j in range(2, C):                   # 2열 이상
        if arr[i][j] > 4:                       # 문제조건상 4이하면 실질적으로 확산이 안됨
            spread = arr[i][j] // 5
            if i < R - 1:                           # down
                arr_tmp[i + 1][j] += spread
                arr[i][j] -= spread
            if i > 0:                               # up
                arr_tmp[i - 1][j] += spread
                arr[i][j] -= spread
            if j < C - 1:                           # right
                arr_tmp[i][j + 1] += spread
                arr[i][j] -= spread
            arr_tmp[i][j - 1] += spread             # left
            arr[i][j] -= spread
        arr_tmp[i][j] += arr[i][j]

2.4. 파이썬의 슬라이싱 활용하기

파이썬의 빠른 슬라이싱 속도는 파이썬의 장점 중 하나라고 한다.

이를 공기청정기로 인한 공기 순환 코드에 사용할 수 있다.

가로 방향의 공기 순환 시에, for문 대신에 슬라이싱을 이용하면 한 번에 많은 값을 업데이트 할 수 있다.

 

실제로 for문을 슬라이싱으로 바꿔서만 돌렸을 때, 약간의 시간 개선이 있었다.

# 가장 위쪽 0행에서의 슬라이싱 활용 예시

# for문 사용
for j in range(C-1):
	arr[0][j] = arr[0][j+1]
    
# 슬라이싱 사용
arr[0][:C-1] = arr[0][1:C]

# 공기 청정
# Counter_ClockWise
for i in range(A1 - 1, 0, -1):
    arr[i][0] = arr[i-1][0]
# for j in range(C-1):
#     arr[0][j] = arr[0][j+1]
arr[0][:C-1] = arr[0][1:C]
for i in range(A1):
    arr[i][C-1] = arr[i+1][C-1]
# for j in range(C-1, 1, -1):
#     arr[A1][j] = arr[A1][j-1]
arr[A1][2:C] = arr[A1][1:C - 1]
arr[A1][1] = 0

# ClockWise
for i in range(A2 + 1, R-1):
    arr[i][0] = arr[i+1][0]
# for j in range(C-1):
#     arr[R-1][j] = arr[R-1][j+1]
arr[R-1][:C-1] = arr[R-1][1:C]
for i in range(R-1, A2, -1):
    arr[i][C-1] = arr[i-1][C-1]
# for j in range(C-1, 1, -1):
#     arr[A2][j] = arr[A2][j-1]
arr[A2][2:C] = arr[A2][1:C - 1]
arr[A2][1] = 0

---

3. 코드

import sys
input = sys.stdin.readline

R, C, T = map(int, input().split())
arr = [list(map(int, input().split())) for _ in range(R)]

A1 = -1
A2 = -1
for i in range(R):
    if arr[i][0] == -1:
        A1 = i
        A2 = i + 1
        break

for _ in range(T):
    # 확산
    arr_tmp = [[0] * C for _ in range(R)]
    for i in range(R):
        for j in range(2):                      # 0, 1 열
            if arr[i][j] > 4:                       # 문제조건상 4이하면 실질적으로 확산이 안됨
                spread = arr[i][j] // 5
                if i < R - 1 and arr[i+1][j] != -1:     # down
                    arr_tmp[i + 1][j] += spread
                    arr[i][j] -= spread
                if i > 0 and arr[i-1][j] != -1:         # up
                    arr_tmp[i - 1][j] += spread
                    arr[i][j] -= spread
                if j < C - 1:                           # right
                    arr_tmp[i][j + 1] += spread
                    arr[i][j] -= spread
                if j > 0 and arr[i][j-1] != -1:         # left
                    arr_tmp[i][j - 1] += spread
                    arr[i][j] -= spread
            arr_tmp[i][j] += arr[i][j]
        for j in range(2, C):                   # 2열 이상
            if arr[i][j] > 4:                       # 문제조건상 4이하면 실질적으로 확산이 안됨
                spread = arr[i][j] // 5
                if i < R - 1:                           # down
                    arr_tmp[i + 1][j] += spread
                    arr[i][j] -= spread
                if i > 0:                               # up
                    arr_tmp[i - 1][j] += spread
                    arr[i][j] -= spread
                if j < C - 1:                           # right
                    arr_tmp[i][j + 1] += spread
                    arr[i][j] -= spread
                arr_tmp[i][j - 1] += spread             # left
                arr[i][j] -= spread
            arr_tmp[i][j] += arr[i][j]

    arr = arr_tmp

    # 공기 청정
    # Counter_ClockWise
    for i in range(A1 - 1, 0, -1):
        arr[i][0] = arr[i-1][0]
    arr[0][:C-1] = arr[0][1:C]
    for i in range(A1):
        arr[i][C-1] = arr[i+1][C-1]
    arr[A1][2:C] = arr[A1][1:C - 1]
    arr[A1][1] = 0

    # ClockWise
    for i in range(A2 + 1, R-1):
        arr[i][0] = arr[i+1][0]
    arr[R-1][:C-1] = arr[R-1][1:C]
    for i in range(R-1, A2, -1):
        arr[i][C-1] = arr[i-1][C-1]
    arr[A2][2:C] = arr[A2][1:C - 1]
    arr[A2][1] = 0

result = 0
for row in arr:
    result += sum(row)
print(result + 2)