[Codeforces Round #966] E - Photoshoot for Gorillas

 

 

 

https://codeforces.com/contest/2000/problem/E

알고리즘 : 그리디

 

1. 문제요약

• n * m 격자로 표현되는 정글에 w마리의 고릴라들이 있다. 한 격자에는 최대 한마리까지 들어갈 수 있다. 
• 고릴라들의 키가 주어진다. 키를 고려하여 고릴라를 정글 격자에 적절히 배치시킨다.
• 정글에 포함되는 모든 k*k 정사각형 부분마다 고릴라들의 키의 합을 구한다. (k는 n과 m 이하의 값이다)
• 위에서 구한 값들의 합을 최대화하도록 고릴라를 배치할 때, 그 최대값을 출력한다.
• 테스트케이스의 개수 t가 주어지고, 각 테스트케이스마다 첫번째 줄에 n,m,k의 값을, 두번째 줄에 w의 값을, 마지막 줄에 고릴라의 키 배열이 입력으로 주어진다. 입출력 예시는 아래와 같다.

2. 아이디어

• 모든 정글 격자에 고릴라를 배치할 수 는 없으므로 적절하게 배치할 필요가 있다. 정글 격자의 입장에서 생각해 보면 가장 많은 k*k 정사각형에 포함되는 격자에 키가 큰 고릴라를 배치하는게 유리함을 알 수 있다.
• 각 정글 격자마다 자신을 포함하는 정사각형의 개수를 체크하여 n*m 크기의 배열에 저장한 후 역순으로 정렬시킨다.
• 고릴라의 키를 역순으로 정렬시키고, 위 배열과 매칭시켜 계산하면 구하는 최대값을 얻을 수 있다.

3. 느낀점

• 무난하게 풀긴 하였으나, n*m 배열을 만들기 전 단계에서 크기 x의 1차원 배열에 크기 k만큼의 1들을 처음부터 끝까지 더한 배열을 얻는 함수를 만드는 아이디어를 떠올리기가 은근히 어려웠다.
• 위 아이디어는 부분합에 기반한 것인데, 여전히 부분합을 자유자재로 쓰고 있지 못하고 있음을 느꼈다.
• 아래 코드에서 myarr 함수 부분은 기억해 두면 좋을 것 같다.

4. 파이썬 코드

def myarr(x, k):
    arr = [0] * x
    for i in range(x - k + 1):
        arr[i] += 1
        if i + k < x:
            arr[i + k] -= 1
    for i in range(1, x):
        arr[i] += arr[i - 1]
    return arr


import sys
input = sys.stdin.readline

t = int(input())
for _ in range(t):
    n, m, k = map(int, input().split())
    w = int(input())
    alist = list(map(int, input().split()))
    alist.sort(reverse=True)

    nlist = myarr(n, k)
    mlist = myarr(m, k)

    nmlist = []
    for i in range(n):
        for j in range(m):
            nmlist.append(nlist[i] * mlist[j])
    nmlist.sort(reverse=True)

    ans = 0
    for a, b in zip(alist, nmlist):
        ans += a * b
    print(ans)