두 개의 배열 A와 B가 있다고 가정해 보겠습니다. 배열 A의 크기는 행(row)의 개수를 의미하며, A[i]는 i번째 행에 있는 상자의 개수를 나타냅니다. 배열 B는 공(ball)들의 목록으로, B[i]는 각 공에 적힌 숫자입니다. 이때 i번째 공(값이 B[i])은 처음 위치부터 세었을 때 B[i]번째에 해당하는 상자에 놓이게 됩니다. 따라서 우리가 구해야 할 것은 각 B[i]에 대응하는 상자의 행과 열입니다.
문제 예시
입력이 A = [3, 4, 5, 6], B = [1, 3, 5, 2]라고 한다면, 출력은 [(1, 1), (1, 3), (2, 2), (1, 2)]가 됩니다.
- B[0] = 1 → 1행 1열
- B[1] = 3 → 1행 3열
- B[2] = 5 → 2행 2열
- B[3] = 2 → 1행 2열
첫 번째 행에는 상자가 3개 있으므로, 4번째 상자부터는 두 번째 행에 속하게 됩니다. 즉, 값 5인 공은 전체 상자 중 5번째 자리이므로 2행 2열에 위치하는 것입니다.
해결 접근 방법
이 문제는 누적 합(prefix sum)과 이진 탐색(binary search)을 활용하면 효율적으로 해결할 수 있습니다. 단계별로 살펴보겠습니다.
- 배열 A를 누적 합 형태로 변환합니다. 변환 후 A[i]는 0번째 행부터 i번째 행까지의 상자 총 개수를 의미합니다.
- 각 공의 값 B[i]에 대해
bisect_left를 사용해 B[i]가 삽입될 수 있는 인덱스를 찾습니다. 이 인덱스가 곧 공이 속한 행이 됩니다. - 찾은 인덱스가 1 이상이라면, B[i]에서 바로 이전 행까지의 누적 합(A[row - 1])을 빼서 해당 행 내에서의 열 위치(box_num)를 계산합니다.
- 인덱스가 0이라면 첫 번째 행에 속하는 것이므로 box_num은 B[i] 그대로입니다.
- 마지막으로 (row + 1, box_num) 형태의 좌표를 출력합니다. 인덱스는 0부터 시작하지만 실제 행·열 번호는 1부터 시작하기 때문입니다.
구현 코드
import bisect
def get_position(A, B):
len_a = len(A)
len_b = len(B)
# 누적 합으로 변환
for i in range(1, len_a):
A[i] += A[i - 1]
# 각 공의 위치 찾기
for i in range(len_b):
row = bisect.bisect_left(A, B[i])
if row >= 1:
box_num = B[i] - A[row - 1]
else:
box_num = B[i]
print((row + 1, box_num))
A = [3, 4, 5, 6]
B = [1, 3, 5, 2]
get_position(A, B)
입력
[3, 4, 5, 6], [1, 3, 5, 2]
출력
(1, 1)
(1, 3)
(2, 2)
(1, 2)
복잡도 분석
누적 합 변환에는 O(n)의 시간이 걸리며(n은 행의 개수), 각 공의 위치를 찾을 때 이진 탐색을 사용하므로 공 하나당 O(log n)이 소요됩니다. 따라서 m개의 공에 대한 전체 시간 복잡도는 O(m log n + n)입니다. 만약 매번 선형 탐색으로 행을 찾았다면 O(m × n)이 되어 비효율적일 수 있으므로, bisect 모듈을 활용한 이진 탐색이 성능 면에서 큰 이점을 제공합니다.