[SWEA 2115] 벌꿀 채취

접근

두 명의 일꾼이 각각 가로로 연속으로 M개의 벌통을 선택한다. 이때 두 명의 일꾼이 선택한 벌통이 겹치면 안된다. 

=> 이는 2차원 리스트의 완전탐색을 통해 두 명의 일꾼이 벌통을 선택하는 경우의 수를 구할 수 있을 것이다.

 

선택한 M개의 벌통 중에서 가중치가 가장 큰 것을 찾아야한다. 예를 들어 5, 5, 9 벌통을 선택했고 C가 10이라고 하자. 그렇다면 최대 수익을 위해 5, 5가 아닌 9를 선택해야한다. 

=> itertools의 combinations를 활용하여 선택한 벌통 중에 최대수익을 찾는다.

 

구현

- 벌통을 2개 선택했을 때 최대수익을 구한다.

- 두 명의 일꾼이 선택한 벌통(storage[0], storage[1])에서 combinations를 통해 최대수익을 구한다.

- 이를 result에 더하여 answer와 최댓값 비교를 진행한다.

- 위의 조건이 아닐 시에는 벌통을 선택한다.

- 현재 위치에서 벌통을 선택할 수 있다면 (범위를 벗어나지 않고, 다른 일꾼과 겹치지않는다면) 벌통을 선택하도록 한다.

def solve(visited, picked):
    global answer
    if picked == 2:
        result = 0
        for i in range(2):
            total = 0
            for j in range(1, M+1):
                for combi in combinations(storage[i], j):
                    if sum(combi) <= C:
                        temp = 0
                        for c in combi:
                            temp += (c * c)
                        total = max(total, temp)
            result += total
        answer = max(answer, result)
        return


    for y in range(N):
        for x in range(N):
            if is_possible(y, x, visited):
                for d in range(M):
                    ny = y
                    nx = x + (d * 1)
                    visited[ny][nx] = True
                    storage[picked].append(board[ny][nx])
                solve(visited, picked+1)

                for d in range(M):
                    ny = y
                    nx = x + (d * 1)
                    visited[ny][nx] = False
                    storage[picked].pop()

---

전체 코드

from itertools import combinations
testCase = int(input())

def check_range(x):
    return (0 <= x < N)

def is_possible(y, x, visited):
    for d in range(M):
        ny = y
        nx = x + (d * 1)
        if check_range(nx) and visited[ny][nx] == False:
            continue
        return False
    return True



def solve(visited, picked):
    global answer
    if picked == 2:
        result = 0
        for i in range(2):
            total = 0
            for j in range(1, M+1):
                for combi in combinations(storage[i], j):
                    if sum(combi) <= C:
                        temp = 0
                        for c in combi:
                            temp += (c * c)
                        total = max(total, temp)
            result += total
        answer = max(answer, result)
        return


    for y in range(N):
        for x in range(N):
            if is_possible(y, x, visited):
                for d in range(M):
                    ny = y
                    nx = x + (d * 1)
                    visited[ny][nx] = True
                    storage[picked].append(board[ny][nx])
                solve(visited, picked+1)

                for d in range(M):
                    ny = y
                    nx = x + (d * 1)
                    visited[ny][nx] = False
                    storage[picked].pop()

for tc in range(testCase):
    N, M, C = map(int, input().split())
    board = [list(map(int, input().split())) for _ in range(N)]
    visited = [[False for _ in range(N)] for _ in range(N)]
    storage = [[] for _ in range(2)]
    answer = 0
    solve(visited, 0)
    print("#{} {}".format(tc+1, answer))