[백준 28216] 2023 정올 초등부 1차 아이템 획득

문제 출처 : https://www.acmicpc.net/problem/28216

28216번: 아이템 획득

이 문제는 2023년도 정보 올림피아드 1차 초등부 3번 문제로 출제 되었던 문제 입니다.

문제 이해하기

2차원 지도에서 아이템을 모으는 문제로 인접 행렬로 구현하면 될 것 같은 문제 입니다. 하지만 지도의 크기가 20만 입니다. 또한 이동하는 횟수인 Q 역시 20만 입니다. 쉽게 시간복잡도가 20만 * 20만이라는 것을 알 수 있고 시간초과가 예상되기 때문에 인접 행렬로 문제를 풀면 안될 것 같습니다. 도저히 방법을 모르겠다면은 지도를 하나하나 탐색하여 부분점수를 얻는 것도 한 방법 입니다. 하지만 100점을 얻기 위해서는 다른 방법을 생각해야 합니다.

이런 문제에서 가장 먼저 생각할 수 있는 방법이 바로 누적합 입니다. 시작 지점을 알고 이동하는 거리를 알고 있기 때문에 끝의 위치에서 시작 위치의 누적된 아이템을 가져오면 하나하나 확인하지 않아도 문제를 해결할 수 있습니다.

하지만 여기에 또 하나의 문제가 있습니다. 바로 지도가 너무 크다는 것입니다. 인접 행렬 20만 * 20만의 배열이기 때문에 x축, y축 각각 누적합을 만드는데 너무 많은 메모리 소모가 됩니다. 따라서 좌표를 압축하여 그 위치에 대한 누적합을 만들어야 메모리 낭비를 줄일 수 있습니다. 즉 우리가 만들어야할 리스트는 압축된 좌표를 가지고 있는 리스트와 누적합을 가지고 있는 리스트가 있어야 합니다.

배열 구성하기

예제를 2차원 배열로 나타낸 것입니다. 딱 봐도 메모리의 낭비가 심하다는 것을 알 수 있습니다. 이것을 다음과 같은 형태로 바꿔주는 것입니다. xp는 x의 위치를 나타냅니다. yp는 y의 위치를 나타냅니다. xt는 x축의 누적합을 나타내고, yt는 y축의 누적합을 나타낸 것입니다.

• xp[1] = [5], xt[1] = [1]

x의 첫 번째 라인에는 하나의 상자가 있고 그 위치는 5라는 뜻입니다. 그것이 xp[1] = 5 입니다. 다음으로 그 위치에 있는 상자의 아이템 개수는 1개 입니다. 그것이 xt[1] = [1] 입니다.

• xp[3] = [5], xt[3] = [2]

x의 세 번째 라인에 상자의 위치는 5 입니다. 그리고 그 상자에는 2개의 아이템이 있다는 뜻입니다.

• xp[5] = [5, 8], xt[5] = [3, 8]

x의 5번째 라인에는 2개의 상자가 있고 각각의 위치가 5, 8 입니다. 그리고 상자에는 3과 5개의 아이템이 있어 누적합으로 3, 8이 있습니다.

• yp[5] = [1, 3, 5], yt[5] = [1, 3, 6]

y의 5번째 라인에는 3개의 상자가 있고 그 위치는 1, 3, 5 입니다. 그리고 상자에는 아이템이 1, 2, 3개씩 존재하기 때문에 누적합은 1, 3, 6이 됩니다.

• yp[8] = [5], yt[8] = [5]

마지막으로 y의 8번째 라인에는 5번 위치에 상자가 있고 그 상자의 아이템 개수는 5개 입니다.

이것으로 2차원 지도의 모든 정보를 리스트에 넣었습니다. 이제 이것을 통해 어떻게 아이템의 개수를 얻는지 알아보겠습니다.

아이템 얻기

첫 번째 이동 0, 4를 생각해 보겠습니다. 자동차의 처음 위치는 1, 1 입니다. 여기서 x축으로 4만큼 이동합니다. 그럼 x축으로 5까지 이동하게 됩니다. xp[5]을 보니 5와 8 위치에 아이템이 있습니다. 현재 y위치는 1이기 때문에 아무 아이템을 얻지 못합니다.

다음으로 1, 9 입니다. 1은 y축 방향으로 이동하라는 뜻이고, 9는 위치를 9만큼 이동하라는 뜻입니다.

xp[5]를 보니 5, 8이 있습니다. 우리는 이 위치를 모두 지나게 됩니다. 따라서 xt[5]에 있는 누적합인 8을 얻게 됩니다.

다음으로 3, 5 입니다. 3은 y축이 감소하는 방향이고, 5는 5칸 이동하라는 뜻입니다.

10에서 부터 5까지 이동하기 때문에 먼저 xp[5]를 확인합니다. 5, 8 모두 경유하기 때문에 이번에도 xt[5]에 있는 누적합 8을 모두 얻습니다. 이것으로 지금까지 얻은 아이템은 8 + 8로 16이 됩니다.

다음으로 2, 3 입니다. 2는 x축 감소하는 방향으로 3만큼 이동하라고 합니다.

yp[5]에서 위치를 확인해보면 5에서 2로 이동하였습니다. yp[5] = [1, 3, 5] 로 아이템을 가지고 있습니다. 중요한 점은 시작 지점의 아이템은 얻을 수 없습니다. 따라서 xt[5] = [1, 3, 6]에서 3의 위치에 있는 누적합은 3이고, 1의 위치는 얻지 않기 때문에 3 - 1 로 2를 얻습니다. 이렇게 이동을 하며 최종적으로 얻는 아이템의 개수는 24개가 됩니다.

코드 작성하기

그럼 직접 코드를 작성하면서 이해해 보도록 하겠습니다.

지도 입력받기

mii = lambda : map(int, input().split())
N, Q = mii()
T = 200_001
xa = [[] for _ in range(T)]
ya = [[] for _ in range(T)]

for _ in range(N):
    x, y, w = mii()
    xa[x].append((y, w))
    ya[y].append((x, w))

20만 * 20만 크기의 2차원 배열의 지도를 만들면 메모리 초과의 위험이 있기 때문에 x축, y축을 각각 xa, ya로 만들었습니다. 입력 받은 위치와 아이템 w값으로 리스트를 만들어 주었습니다.

위치 리스트, 누적합 리스트 만들기

def make_list(arr, pos, total):
    arr.sort()
    t = 0
    for p, items in arr:
        pos.append(p)
        t += items
        total.append(t)

xp = [[] for _ in range(T)]
yp = [[] for _ in range(T)]
xt = [[] for _ in range(T)]
yt = [[] for _ in range(T)]

for i in range(T):
    make_list(xa[i], xp[i], xt[i])
    make_list(ya[i], yp[i], yt[i])

xa, ya 리스트를 통해 각각의 위치를 정렬하여 나타낸 xp 리스트를 만들고 xp, yp의 위치와 대응하는 인덱스에 아이템 리스트를 누적합을 가진 xt, yt 리스트를 만들었습니다. t값에 계속 items의 값을 누적하여 만들었습니다. 누적합을 만드는 방법에 대해서는 아래 링크에서 자세히 다루었으니 확인해 보시기 바랍니다.

https://wikidocs.net/206294

02. 누적 합

이동 로직 구현

x, y, ans = 1, 1, 0
for _ in range(Q):
    d, v = mii()
    if d == 0:
        ans += get_items(yp[y], yt[y], x + 1, x + v)
        x += v
    elif d == 1:
        ans += get_items(xp[x], xt[x], y + 1, y + v)
        y += v
    elif d == 2:
        ans += get_items(yp[y], yt[y], x - v, x - 1)
        x -= v
    elif d == 3:
        ans += get_items(xp[x], xt[x], y - v, y - 1)
        y -= v

print(ans)

첫 시작은 1, 1 입니다. 입력받은 d 값에 따라 x축, y축으로 이동하는 로직을 만들었습니다. 각각의 정보를 입력하여 get_items 함수를 통해 아이템의 개수를 리턴 받게 됩니다.

get_items 함수 구현

from bisect import bisect_left, bisect_right

def get_items(pos, total, left, right):
    start = bisect_left(pos, left)
    end = bisect_right(pos, right) - 1

    result = 0
    if end < start:
        result = 0
    else:
        result = total[end]
        if 0 < start:
            result -= total[start - 1]

    return result

get_items 함수 입니다. pos는 x축을 이동할 경우 yp리스트를, y축을 이동할 경우 xp리스트를 이동합니다. total은 이동한 리스트의 인덱스와 일치하는 누적합 리스트 xt, yt를 받습니다. 그리고 누적합을 구할 left, right 값을 입력 받습니다.

left와 right 값을 통해 리스트에서 인덱스를 찾아야 합니다. xp, yp 값은 이미 정렬되어 있기 때문에 이분탐색을 사용하면 쉽게 해당 인덱스를 찾을 수 있습니다. 이분 탐색 함수인 bisect_left와 bisect_right에 대해서는 아래 링크에 설명 되어 있습니다.

https://wikidocs.net/206313

05. 이분 탐색

bisect_right를 사용할 경우 인덱스를 벗어나는 경우가 있습니다. 또는 리스트에 아무것도 없는 경우가 있을 수 있습니다. 이런 경우를 대비하여 리스트 인덱스 범위를 잘 파악해야 합니다. 범위를 구하다보면 시작보다 끝이 작은 경우도 있고, 시작이 0보다 작은 경우도 있을 수 있습니다. 이런 경우를 잘 생각해서 누적합을 구해야 합니다.

전체 코드

그럼 전체 코드를 확인해 보겠습니다.

from bisect import bisect_left, bisect_right

mii = lambda : map(int, input().split())
N, Q = mii()
T = 200_001
xa = [[] for _ in range(T)]
ya = [[] for _ in range(T)]

for _ in range(N):
    x, y, w = mii()
    xa[x].append((y, w))
    ya[y].append((x, w))

def make_list(arr, pos, total):
    arr.sort()
    t = 0
    for p, items in arr:
        pos.append(p)
        t += items
        total.append(t)

xp = [[] for _ in range(T)]
yp = [[] for _ in range(T)]
xt = [[] for _ in range(T)]
yt = [[] for _ in range(T)]

for i in range(T):
    make_list(xa[i], xp[i], xt[i])
    make_list(ya[i], yp[i], yt[i])

def get_items(pos, total, left, right):
    start = bisect_left(pos, left)
    end = bisect_right(pos, right) - 1

    result = 0
    if end < start:
        result = 0
    else:
        result = total[end]
        if 0 < start:
            result -= total[start - 1]

    return result

x, y, ans = 1, 1, 0
for _ in range(Q):
    d, v = mii()
    if d == 0:
        ans += get_items(yp[y], yt[y], x + 1, x + v)
        x += v
    elif d == 1:
        ans += get_items(xp[x], xt[x], y + 1, y + v)
        y += v
    elif d == 2:
        ans += get_items(yp[y], yt[y], x - v, x - 1)
        x -= v
    elif d == 3:
        ans += get_items(xp[x], xt[x], y - v, y - 1)
        y -= v

print(ans)