백준 23291 어항 정리 (파이썬)

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

23291번: 어항 정리

1. 어항을 쌓는 식으로 구현하기 위해서 빈 리스트를 원소로 갖는 길이 n의 리스트 fish를 만든다

   fish에 물고기 초기값을 각 어항에 넣어준다

2. 조건을 만족할때까지 while문으로 반복한다

   우선 최소값을 갖는 모든 어항에 1을 더한다

3. 어항 맨 왼쪽칸을 다음 칸 위에 append 한다

   마지막 어항부터 시작해서, 해당 어항의 길이가 1보다 크면, 회전 후 다음칸에 쌓는다

   이 때 어항의 길이와 바닥에 위치할 어항의 길이를 비교해서 쌓을 수 있는지 검사한다.

   쌓을 수 있다면, 반복문에 zip을 사용해서 현재 어항 오른쪽 칸부터 순서대로 쌓아준다

4. 쌓는게 끝나면 fish_move 함수로 물고기 수를 조절한다

   모든 칸을 동시에 처리해야 하기때문에 fish_temp에 fish를 복사하고, fish_temp 수를 조절한다

   문제에서 어항 쌓은 그림 기준 왼쪽 하단부터 오른칸, 윗칸을 비교해서 수를 조절해준다

   조절이 끝나면 fish를 fish_temp로 바꿔준다

5. make_line 함수로 일렬로 만든다

   왼쪽 어항부터 시작해서 물고기가 있으면 temp에 해당 어항을 복사하고 비워준다. (중복 처리 방지)

   temp에 있는 값을 순서대로 불러와서 blank_idx에 해당하는 어항에 채워주면 된다

6. fish를 left, right 리스트로 나눠준다

   left는 일렬이므로 역순으로 바꿔주면 180도 회전이 된다

   right에 역순으로 바꾼 left를 하나씩 쌓아준다

7. right를 다시 left, right로 나눠준다

   left[::-1]에 zip(*)를 두번 하면 180도 회전이 된다

   right에 회전한 left를 하나씩 쌓아준다

8. right를 fish_move에 입력해서 수를 조절한다

9. make_line 함수에 맞추기 위해서 빈 리스트 fish를 다시 만들고 오른쪽부터 채워준 후 make_line을 실행한다

10. 최대값, 최소값의 차가 k보다 크지않으면 cnt를 출력하고 전체 반복문을 종료한다

import sys
from copy import deepcopy

input = sys.stdin.readline

n, k = map(int, input().split())
a = list(map(int, input().split()))

fish = [[] for _ in range(n)]
for i in range(n):
    fish[i].append(a[i])


def fish_move(n, fish):
    fish_temp = deepcopy(fish)
    for i in range(n):
        if fish[i]:
            for j in range(len(fish[i])):
                if i + 1 < len(fish):
                    if j < len(fish[i + 1]):
                        if fish[i + 1][j]:
                            if fish[i][j] > fish[i + 1][j]:
                                diff = (fish[i][j] - fish[i + 1][j]) // 5
                                fish_temp[i][j] -= diff
                                fish_temp[i + 1][j] += diff
                            elif fish[i][j] < fish[i + 1][j]:
                                diff = (fish[i + 1][j] - fish[i][j]) // 5
                                fish_temp[i][j] += diff
                                fish_temp[i + 1][j] -= diff

                if j + 1 < len(fish[i]):
                    if fish[i][j] > fish[i][j + 1]:
                        diff = (fish[i][j] - fish[i][j + 1]) // 5
                        fish_temp[i][j] -= diff
                        fish_temp[i][j + 1] += diff
                    elif fish[i][j] < fish[i][j + 1]:
                        diff = (fish[i][j + 1] - fish[i][j]) // 5
                        fish_temp[i][j] += diff
                        fish_temp[i][j + 1] -= diff
    return fish_temp


def make_line():
    blank_idx = 0
    for i in range(n):
        if len(fish[i]) > 1:
            temp = fish[i]
            fish[i] = []
            for k in temp:
                fish[blank_idx].append(k)
                blank_idx += 1


cnt = 1
while True:
    min_n = sys.maxsize
    for i in range(len(fish)):
        min_n = min(fish[i][0], min_n)
    for i in range(len(fish)):
        if fish[i][0] == min_n:
            fish[i][0] += 1

    fish[1].append(fish[0][0])
    fish[0] = []

    flag = 0
    while True:
        for i in range(len(fish)-1, 0, -1):
            if len(fish[i]) > 1:
                if len(fish[i]) > len(fish) - i - 1:
                    flag = 1
                    break
                for j in range(i + 1, n):
                    if fish[j]:
                        for v, idx in zip(fish[i], range(j, j + len(fish[i]) + 1)):
                            fish[idx].append(v)
                        break
                fish[i] = []
        if flag == 1:
            break

    fish = fish_move(n, fish)

    make_line()

    left, right = fish[:n//2], fish[n//2:]
    n_left = left[-1::-1]
    for l, r in zip(n_left, right):
        r.append(l[0])

    left, right = right[:n//4], right[n//4:]
    for _ in range(2):
        left = list(zip(*left[::-1]))
    for l, r in zip(left, right):
        for ll in l:
            r.append(ll)

    right = fish_move(n//4, right)

    fish = [[] for _ in range(n)]
    for i in range(len(right)):
        fish[-i-1] = right[len(right) - i -1]
    make_line()

    max_n, min_n = 0, sys.maxsize
    for f in fish:
        max_n = max(f[0], max_n)
        min_n = min(f[0], min_n)

    if max_n - min_n <= k:
        print(cnt)
        break

    cnt += 1