[백준] 18513 샘터 python

https://www.acmicpc.net/problem/18513

18513번: 샘터

그래프 탐색인데 1차원이라 생소하긴했지만 어려운 문제는 아니다.

 

단지 이전부터 많이 해왔던 것처럼 visit을 처음부터 잡아놓고 진행하면 메모리 초과가 나온다. 범위가 +-100,000,000이기 때문에 4 * 200,000,000 이므로 약 762MB 쯤 .

따라서 set이나 dictionary 자료구조를 사용해서 visit을 체크하면 된다.

 

import sys; input = sys.stdin.readline
from collections import deque

def bfs(q, visit, limit):
    unhappiness = 0
    ind = 0
    flag = False
    while q:
        cur = q.popleft()
        for d in [-1, 1]:
            next = cur + d
            if visit.get(next) is None:
                visit[next] = visit[cur] + 1
                q.append(next)
                unhappiness += visit[next]
                ind += 1
                if ind == limit:
                    flag = True
                    break
        if flag:
            break

    return unhappiness


def main():
    n, k = map(int, input().split())
    visit = {}
    q = deque()
    for i in map(int, input().split()):
        visit[i] = 0
        q.append(i)

    print(bfs(q, visit, k))


if __name__ == "__main__":
    main()