[백준] 2586 소방차 Python 3

문제 설명

이 문제는 여러 대의 소방차가 펌프에서 물을 채우기 위해 호스를 연결할 때, 호스의 총 길이를 최소화하는 문제입니다. 이 문제를 해결하기 위해 펌프와 소방차의 위치를 고려하여 가장 가까운 펌프와 소방차를 연결하는 방법을 구현했습니다. 이 코드는 bisect 모듈을 사용하여 이진 탐색을 통해 가장 가까운 펌프를 찾고, 동적 계획법을 사용하여 이미 사용된 펌프를 피하는 로직을 포함하고 있습니다.

upper_bound 함수

upper_bound 함수는 주어진 배열에서 값이 들어갈 위치를 찾는 함수입니다. 이 함수는 bisect 모듈의 bisect_right 함수를 사용하여 배열에서 value보다 큰 값이 처음 나타나는 위치를 반환합니다. 이 함수는 소방차의 위치에 가장 가까운 펌프를 찾는 데 사용됩니다.

solve 함수

이 함수는 주어진 펌프와 소방차의 위치를 기반으로, 총 호스 길이를 최소화하는 로직을 구현합니다. 각 소방차에 대해 다음 단계를 수행합니다:

1. upper_bound를 사용하여 소방차 위치를 기준으로 가장 가까운 펌프를 찾습니다.
2. 가장 가까운 펌프가 아직 사용되지 않았다면, 해당 펌프를 사용하고 호스 길이를 결과에 추가합니다.
3. 이미 사용된 펌프라면, 왼쪽과 오른쪽을 탐색하여 사용 가능한 펌프를 찾습니다.
4. 추가적인 계산을 통해, 가장 적합한 펌프를 선택하고 호스 길이를 결과에 추가합니다.

문제 해결

이 문제를 해결하기 위해 파이썬에서 제공하는 bisect 모듈을 사용하여 이진 탐색을 구현했습니다. 이진 탐색을 통해 각 소방차에 대해 가장 가까운 펌프를 찾아내고, 이미 사용된 펌프를 피하기 위해 추가적인 로직을 구현하였습니다.

다음은 문제를 해결하기 위한 코드입니다.

 

import sys
import bisect

# upper_bound 함수: 주어진 배열에서 값이 들어갈 위치를 찾는 함수
def upper_bound(array, value):
    index = bisect.bisect_right(array, value)
    return index

# 문제를 해결하는 메인 함수
def solve(n, m, p, f):
    dp = [False] * (n + 10)  # 펌프 사용 여부를 저장하는 배열
    res = 0  # 결과값 (총 호스 길이)
    
    for i in range(1, m + 1):
        ub = upper_bound(p, f[i])  # 현재 소방차 위치를 기준으로 upper bound를 찾음
        # 가까운 펌프를 선택하는 로직
        node = ub - 1 if abs(p[ub] - f[i]) > abs(p[ub - 1] - f[i]) else ub
        
        if not dp[node]:  # 선택한 펌프가 아직 사용되지 않았다면
            dp[node] = True
            res += abs(p[node] - f[i])  # 호스 길이를 추가
        else:  # 선택한 펌프가 이미 사용되었다면
            cnt = 0
            for j in range(node, 0, -1):  # 왼쪽으로 빈 펌프를 찾음
                if dp[j]:
                    cnt += 1
                else:
                    break
            
            if node - cnt == 0:  # 왼쪽으로 더 이상 빈 펌프가 없으면
                cnt = 0
                for j in range(node, m + 1):  # 오른쪽으로 빈 펌프를 찾음
                    if dp[j]:
                        cnt += 1
                    else:
                        break
                dp[node + cnt] = True
                res += abs(p[node + cnt] - f[i])  # 호스 길이를 추가
            else:  # 왼쪽에 빈 펌프가 있으면
                add = 0
                push = 0
                ccnt = 0
                for j in range(node + 1, n + 1):
                    if dp[j]:
                        node += 1
                    else:
                        break
                for j in range(node, 0, -1):
                    if dp[j]:
                        add += abs(p[j] - f[i - node + j - 1])
                        push += abs(p[j - 1] - f[i - node + j - 1])
                        ccnt += 1
                    else:
                        break
                push += abs(p[node] - f[i])
                if add + abs(p[node + 1] - f[i]) >= push:
                    res -= add
                    dp[node - ccnt] = True
                    res += push
                else:
                    ccnt = 0
                    for s in range(node, n + 1):
                        if dp[s]:
                            ccnt += 1
                        else:
                            break
                    dp[node + ccnt] = True
                    res += abs(p[node + ccnt] - f[i])
    return res

# 메인 함수
def main():
    input = sys.stdin.read  # 입력을 읽음
    data = input().split()  # 공백을 기준으로 데이터를 분리
    n = int(data[0])  # 펌프의 개수
    m = int(data[1])  # 소방차의 개수
    
    # 펌프 위치 배열
    p = [-2147483647] + [int(data[i]) for i in range(2, n + 2)]
    p += [2147483647] * (100010 - n - 2)
    
    # 소방차 위치 배열
    f = [-2147483647] + [int(data[i]) for i in range(n + 2, n + m + 2)]
    f += [2147483647] * (100010 - m - 2)
    
    result = solve(n, m, p, f)  # 문제를 해결
    print(result)  # 결과를 출력

if __name__ == "__main__":
    main()

결론

이 코드에서는 이진 탐색과 그리디 알고리즘을 활용하여 효율적으로 문제를 해결하였습니다. 소방차의 위치와 펌프의 위치를 정렬하고, 각 소방차가 가장 가까운 펌프에 연결되도록 구현하였습니다.