[프로그래머스] 파이썬 - 석유 시추 (PCCP)

문제

풀이

첫번째 코드 - 시간초과

from collections import deque

def bfs(land,visited,i,j,ground,idx):
    cnt = 1
    dx = [0,0,1,-1]
    dy = [1,-1,0,0]
            
    queue = deque([[i,j]])
    visited[i][j] = True
    
    positions = [[i,j]]
    
    while queue:
        x,y = queue.popleft()

        for i in range(4):
            nx,ny = x+dx[i], y+dy[i]
            
            if nx>=0 and ny>=0 and nx<len(land) and ny<len(land[0]) and not visited[nx][ny] and land[nx][ny] == 1:
                cnt += 1
                visited[nx][ny] = True
                queue.append([nx,ny]) 
                positions.append([nx,ny])

    for pos in positions:
        x,y = pos
        ground[x][y] = [True,idx,cnt]

def solution(land):
    answer = 0
    n = len(land)
    m = len(land[0])
    visited = [[False]*m for _ in range(n)]
    # [석유 존재 여부,?번 석유 땅, 석유 면적] 
    ground = [[[False, 0, 0] for _ in range(m)] for _ in range(n)]
    # 석유 땅 인덱스
    idx = 1
    
    for i in range(n):
        for j in range(m):
            if land[i][j] == 1 and not visited[i][j]:
                bfs(land,visited,i,j,ground,idx)
                idx += 1   

    for i in range(m):
        cnt = 0
        visited_oil = set()
        for j in range(n):
            # 석유가 있는 땅이고 아직 계산하지 않은 석유 땅의 idx인 경우
            if ground[j][i][0] and ground[j][i][1] not in visited_oil:
                visited_oil.add(ground[j][i][1])
                cnt += ground[j][i][2]              
        
        answer = max(answer,cnt)
    
    return answer

 

코드가 길고 비효율적이다. 

처음 풀 때는 ground라는 배열을 하나 선언했다.

land를 bfs로 돌며 석유 땅들을 검사한 후 각 인덱스 별로 석유가 존재하는지 & 존재한다면 몇번째 석유 인지 & 면적은 얼마인지를 계산해 ground에 저장했다.

 

이후 ground를 각 column에 따라 검사하여 답을 도출했다.

정답은 잘 나왔지만, 효율성 테스트에서 시간초과에 걸렸다.

 

두번째 풀이 - 정답 코드

from collections import deque

def bfs(land,visited,i,j,oil):
    cnt = 1
    dx = [0,0,1,-1]
    dy = [1,-1,0,0]
            
    queue = deque([[i,j]])
    visited[i][j] = True
    columns = set()
    columns.add(j)
    
    while queue:
        x,y = queue.popleft()

        for i in range(4):
            nx,ny = x+dx[i], y+dy[i]
            
            if nx>=0 and ny>=0 and nx<len(land) and ny<len(land[0]) and not visited[nx][ny] and land[nx][ny] == 1:
                cnt += 1
                visited[nx][ny] = True
                queue.append([nx,ny]) 
                columns.add(ny)
    
    for c in columns:
        oil[c] += cnt

def solution(land):
    answer = 0
    n = len(land)
    m = len(land[0])
    visited = [[False]*m for _ in range(n)]
    oil = [0] * m 
    
    for i in range(n):
        for j in range(m):
            if land[i][j] == 1 and not visited[i][j]:
                bfs(land,visited,i,j,oil)
                  
    answer = max(oil)
    
    return answer

 

문제를 좀 더 간단히 생각해보았다.

ground라는 배열을 선언할 필요 없이, 그냥 bfs를 계산할 때 해당 석유 면적의 y값들에 면적 값을 더해주면 되는 간단한 문제였다.

 

프로그래머스