백준 10800 : 컬러볼 (Python)

본문 링크

from bisect import bisect_left
import sys
input=sys.stdin.readline

N=int(input())
L=[] # 공과 번호를 입력값으로 담을 배열
color=[ [] for i in range(200001)] ; Perfix_color= [ [0] for _ in range(200001) ]
Perfix_ALL=[0]*(N+1)
ALL_Ball=[]
for i in range(N):
    a,b=map(int,input().split())
    L.append([a,b])
    color[a].append(b)
    ALL_Ball.append(b)

ALL_Ball.sort()
for i in range(1,N+1):
    Perfix_ALL[i]=Perfix_ALL[i-1]+ALL_Ball[i-1]

for i in range(200001):
    if color[i]:
        color[i].sort()
        for j in range(1,len(color[i])+1):
            Perfix_color[i].append(Perfix_color[i][j-1]+color[i][j-1])


for x,y in L:
    total=0
    index_ALL=bisect_left(ALL_Ball,y) # 자기보다 작은 모든 공의 합
    index_color=bisect_left(color[x] , y) # 같은색깔의 자기보다 작은 공의 합
    
    total+=Perfix_ALL[index_ALL]
    total-=Perfix_color[x][index_color]
    print(total)

📌 어떻게 접근할 것인가?

문제를 풀면서 좀 오래 고민했다. 아이디어는 많이 떠올랐지만 문제에서 주어지는 $N$ $(1 ≤ N ≤ 200,000)$ 범위때문에 어떻게 최적화를 할지 고민했다.

누적합의 한가지 특성을 활용하기로 했다.

정렬되어있는 기존 배열과 그 배열의 누적합의 인덱스는 같다.

예를들어보면

• $[3,6,10,15]$

라는 배열이 있다. 이때 누적합 배열은

• $[3,9,19,34]$

이다. 만약 내가 15까지의 누적합을 알고싶다면 어떻게 해야할까?

$Prefix[List.index(15)]$ 를 해주면 된다. 즉 리스트에서 15 의 인덱스 값은 3 이고
누적합의 인덱스 3 이 바로 34가 된다.

즉 특정 값의 위치를 알고있으면 누적합을 바로 찾을 수 있다. 단 이때 기존배열은 정렬되어있는 상태여야 한다.

그리고 자기보다 작은 서로다른 색깔의 공의 합 을 구해야 한다.

따라서 자기보다 작은 모든 공의 합 - 자기와 색깔이 같고 작은 공의 합 을 계산해주면

자기와 색깔이 다르고 작은 공의 합을 구할수 있다.

총 4개의 배열이 필요하다.

• 정렬된 모든 공의 배열 - $ALL$_$Ball$
• 모든 공의 배열의 누적합 배열 - $Perfix$_$ALL$
• 정렬된 색깔별로 나누어진 배열 - $color$
• 색깔마다의 누적합 배열 - $Perfix$_$color$

그리고 정렬된 특정값의 위치를 빠르게 찾아줄 bisect 라이브러리를 사용한다.

자기보다 작은 공의 총 합은
index_ALL=bisect_left(ALL_Ball,y) 을 통해 위치를 구한뒤
Perfix_ALL[index_ALL] 가 총합이고

같은 색깔의 자기보다 작은 공의 총 합은
index_color=bisect_left(color[x] , y) 을 통해 위치를 구한뒤
Perfix_color[x][index_color] 로 그 값을 구한다.

이후 두개의 차이가 색깔이 다른 자기보다 작은 공의 총 합이된다.

✅ 느낀점

한 3~4시간 넘게 고민한거 같았는데 , $N$ $(1 ≤ N ≤ 200,000)$ 범위때문에
공이 하나 주어지면 "누적합 배열을 통해 $O(2)$ 로 무조건 바로 출력해야한다" 라는 생각을 크게 가졌었던것 같다.
근데 다른사람 풀이를 보니 모든 공을 담을 배열을 한번 정렬해주고
$N$ 번의 반복문을 통해서 $0$ 번째 인덱스 부터 자기공의 인덱스 위치까지 누적합을 한뒤
입력받을때 색깔도 함께 입력받아서 같은 색깔이라면 값을 빼준다. 왜냐하면 정렬된 상태에서 자기보다 인덱스가 낮은 공은 모두 자기보다 작은 공이기 때문이다.
이렇게 하면 연산횟수가 아주 크게 늘어날꺼같지만 의외로 내 풀이보다 속도가 아주 약간보다 빠르다.

색깔별로 공을 담는 배열의 크기를 $200001$ 으로 잡고 그것을 정렬한뒤 누적합을 해서
정렬을 $200000$ 번과 누적합 $200000$번을 2번해서 느려진것같다.

원래는 2차원 배열로 풀려했으나 메모리초과가 났기 때문에 어쩔수없이 배열을 여러개 쓰고 분리하는 식으로 했는데 참 접근하기 힘든것같다.

그래도 특정 값의 위치를 알고있으면 누적합을 바로 찾을 수 있다 라는 이론을 이번문제를 통해 알수 있게되었다.