백준 23829 : 인물예술탐사주간 (Python)

본문 링크

import sys
input=sys.stdin.readline
import copy

N,Q=map(int,input().split())
L=sorted(list(map(int,input().split())))

L_sum=copy.deepcopy(L)

for i in range(1, N):
    L_sum[i] = L_sum[i - 1] + L_sum[i]
for i in range(Q):
    index=int(input())

    start=0 ; end=N-1 ; answer=0

    if L[-1]<=index:
        print(N*index - L_sum[-1] )
    elif L[0]>=index:
        print( L_sum[-1] - N*index )
    else:
        while start<=end:
            mid=(start+end)//2
            if L[mid]>=index:
                end=mid-1
            elif L[mid]<index:
                start=mid+1
        print((index*start -L_sum[start-1] ) + ( (L_sum[-1] - L_sum[start-1] ) - index*(N-start)))

📌 어떻게 접근할 것인가?

문제 자체는 쉽다. Q지점에서 사진을 찍었을때 모든 나무의 위치에다가 Q지점을 뺀 절대값을 구하면 된다.

예제입출력을 보면

• 입력
  5 2
  1 3 7 9 10
  4
  12

• 출력
  18
  30

태영이가 $i$ 번째로 사진을 찍은 위치 $X_i$ 보자.
4일때는 (4-1) + (4-3) + (7-4) + (9-4) + (10-4) = 18 이 나오고
12일때는 (12-1) + (12-3) + (12-7) + (12-9) + (12-10) = 30 이 나온다.

이때 유심히 봐야할점은 절대값을 구하는 것이기 때문에 $X_i$ 와 나무의 위치에 따라 서로 뺴는 순서가 바뀐다는 점이다.

그렇다면 단순히 O(N)라고 생각하여 for i in range( len(List) ) 를 선언해주고
total 변수에다가 total+=abs($X_i$-List[i]) 을 해주면되지않을까?

하지만 문제에서 주어지는 범위를 보면

 $1 \leq N,Q \leq 10^5$ 
$1 \leq P_i \leq 10^7$ ($1 \le i \le N$)
$1 \leq X_i \leq 10^7$ ($1 \le i \le Q$)

N,Q는 10만 , $P_i$ 와 $X_i$는 무려 1000만까지 존재한다.
반복문으로 리스트 전체를 탐색하는것은 가능할지라도 저렇게 큰 값을 최대 10만번동안 연산하는건 매우 오래걸린다.

따라서 이때 누적합 알고리즘을 생각해야한다.

문제에서 주어진 리스트를 보자

• 1 3 7 9 10

단순히 누적 합을 적용시키면

• 1 4 11 20 30

위와 같은 형태가 된다.

하지만 아까처럼 시뮬레이션을 직접 해보았을때
4일때를 먼저보면 (4-1) + (4-3) + (7-4) + (9-4) + (10-4) = 18 이 나오고
이는 ( 4x2 ) - (1+3) 과 (7+9+10) - ( 4x3 ) 으로 바라볼수 있다.

여기서 문제는 $X_i$ 위치에 따라 부호가 매번 바뀐다는 점이다.
하지만 $X_i$보다 작은 리스트 값은 $X_i$ x (개수) - ($X_i$ 보다 작은 리스트의 합)이 되고
$X_i$보다 큰 리스트 값은 ($X_i$보다 큰 리스트의 합) - ( $X_i$ x 개수) 이 된다.

그렇다면 우리가 구해야할 값은 딱 하나이다. $X_i$보다 크거나같은 위치는 어디인가?

이것을 완전탐색을 통해 구하면 최대 10만번의 연산이 필요하다.

따라서 더 효율적으로 찾기 위해 이분 탐색 알고리즘을 사용하고자 한다.

📌 어떻게 누적합과 이분탐색 알고리즘을 사용할 것인가?

먼저 import copy 를 통해 누적합 리스트를 copy.deepcopy 로 따로 만들어 준다.

그리고 start= 0 end=N-1로 잡은 후 L[mid] 가 index ( $X_i$ 위치) 보다 크거나 같으면 end 값을 감소 시키고 L[mid] 가 index 보다 작으면 start 값을 증가시킨다.

이때 우리가 찾고자하는 $X_i$ 보다 크거나같은 위치는 start 값이 된다.

이후 다시 기존 배열과 누적합 배열을 보면

• 기존 배열 - 1 3 7 9 10
• 누적 합 배열 -1 4 11 20 30

위와 같고 , $X_i$가 4일때는 start 가 2가 된다. 즉 4보다 큰 인덱스 값은 7이기 때문에 2가 나온다.

$X_i$보다 클때의 계산은 (7-4) + (9-4) + (10-4) -> 26 - 4x3이 되고 여기서 26은
누적합배열의 끝(기존배열의 합) - L[start] -> 30-4=26
그리고 4x3은 $X_i$ x (N-start) -> 4 x (5-2) = 12가 된다
따라서 이 둘을 빼주면 26-12=14가 된다.

다음으로 $X_i$보다 작을때의 계산은 (4-1) + (4-3) -> 4x2 - (1+3) 이 되고 여기서 4x2는
$X_i$ x start 값이 되고
1+3의 값은 누적합의 배열 인덱스에서 start-1값이 된다.
따라서 이 둘을 빼주면 8-4=4가 된다.

마지막으로 14($X_i$보다 클때의 계산값) + 4($X_i$보다 작을때의 계산값) = 18이 성립한다.

이를 한줄로 쓰자면

(index * start -L_sum[start-1] ) + ( (L_sum[-1] - L_sum[start-1] ) - index*(N-start) ) 처럼 쓸수 있다.

여담으로 이러한 문제는 직접 시뮬레이션을 여러번 해보면서 풀이접근을 해보는게 좋다.

✅ 코드에서 주의해야할 점

• 누적합 배열은 deepcopy를 통해 따로 만들어준다.
• 누적합 배열은 리스트를 입력받고 바로 만들어 줘서 한번만에 만들도록한다. 반복문에 넣지말것.
• $X_i$보다 크거나 같은 지점의 최소 인덱스는 start 값을 이용한다.