(Python) 프로그래머스 - 석유 시추

30분이면 풀 수 있을 줄 알았지만 몇 시간을 고민한 문제이다.

각 열을 돌면서 매번 bfs를 탐색해도 시간초과에 안 걸릴 것 같았는데, 효율성 테스트에서 전부 시간 초과가 나서

bfs 탐색이 끝나면 각 열에 오일 개수를 추가하는 로직으로 풀이에 성공했다.

---

로직

일반적인 bfs와 다른점은 bfs 탐색이 끝나면 석유가 있는 열 set을 따로 만들어서 set을 순회하며 원래 oil 리스트에 인접 영역의 석유의 양인 cnt을 더해주는 것이다.

이렇게 하면 bfs 탐색은 한 번만 실행하고 석유가 있는 인접영역을 모두 열에 추가하여 답을 구할 수 있다.

 

최종 코드

from collections import deque


def solution(land):
    answer = 0
    n, m = len(land), len(land[0])

    # 방문배열
    visited = [[False] * m for _ in range(n)]
    # 각 열의 기름의 총량을 저장하는 리스트
    oil = [0] * m
    directions = [(0, 1), (1, 0), (-1, 0), (0, -1)]

    def bfs(row, col):
        queue = deque()
        queue.append([row, col])
        visited[row][col] = True
        cnt = 1
        # bfs 탐색중 석유가 있는 열, 중복을 방지하기 위한 set
        oil_covered = {col}

        while queue:
            pr, pc = queue.popleft()
            for dr, dc in directions:
                nr, nc = pr + dr, pc + dc
                if 0 <= nr < n and 0 <= nc < m and land[nr][nc] == 1 and not visited[nr][nc]:
                    queue.append([nr, nc])
                    visited[nr][nc] = True
                    cnt += 1
                    # 현재 석유를 발견한 열 추가
                    oil_covered.add(nc)

        # 각 열을 돌면서 석유량 추가
        for c in oil_covered:
            oil[c] += cnt

    # bfs 탐색
    for i in range(n):
        for j in range(m):
            if land[i][j] == 1 and not visited[i][j]:
                bfs(i, j)

    answer = max(oil)

    return answer