[백준7578번] 공장 / Python3

문제

어떤 공장에는 2N개의 기계가 2열에 걸쳐 N개씩 배치되어 있다. 이 2개의 열을 각각 A열과 B 열이라고 부른다. A열에 있는 N개의 기계는 각각이 B열에 있는 N개의 기계와 하나씩 짝을 이루어 케이블로 연결되어 있다. 즉, A열의 임의의 기계는 B열의 유일한 기계와 케이블로 연결되어 있고, B열의 임의의 기계는 A열의 유일한 기계와 케이블로 연결되어 있다

또한, 각 기계에는 식별번호가 붙어있으며, 짝이 맺어진 기계끼리는 같은 식별번호가 붙어있다. 즉, 각 열에 있는 N개의 기계끼리는 서로 다른 식별번호를 가지고 있으며, 반대쪽 열에 있는 같은 식별번호를 가진 기계와 케이블로 이어져 있다.

공장 작업의 효율성을 위해 기계들은 짝을 맺은 순서대로 배치되지 않으며, 필요에 따라 각 열의 기계들의 순서를 바꾼 바람에 케이블은 마구 엉켜있는 상태이다. 이렇게 엉켜버린 케이블은 잦은 고장의 원인이 되기 때문에, 기계의 위치를 바꾸지 않은 상태에서 케이블을 두 기계를 잇는 직선의 형태로 만들기로 했다.

예를 들어, 위의 그림과 같이 N = 5이고, A열에 위치한 기계의 식별번호가 순서대로 132, 392, 311, 351, 231이고 B열에 위치한 기계의 식별번호가 순서대로 392, 351, 132, 311, 231이라면 케이블들의 교차 횟수 혹은 서로 교차하는 케이블 쌍의 개수는 3이 된다.

정수 N과 A열에 위치한 기계, B열에 위치한 기계의 식별번호가 각각 순서대로 주어질 때에 서로 교차하는 케이블 쌍의 개수를 정확하게 세어 출력하는 프로그램을 작성하시오.

입력

입력은 세 줄로 이루어져 있다. 첫 줄에는 정수 N이 주어지며, 두 번째 줄에는 A열에 위치한 N개 기계의 서로 다른 식별번호가 순서대로 공백문자로 구분되어 주어진다. 세 번째 줄에는 B열에 위치한 N개의 기계의 식별번호가 순서대로 공백문자로 구분되어 주어진다.

단, 1 ≤ N ≤ 500,000이며, 기계의 식별번호는 모두 0 이상 1,000,000 이하의 정수로 주어진다

출력

여러분은 읽어 들인 2N개의 기계의 배치로부터 서로 교차하는 케이블 쌍의 개수를 정수 형태로 한 줄에 출력해야 한다.

예제 입력 1

5
132 392 311 351 231
392 351 132 311 231

예제 출력 1

3

나의 풀이

세그먼트 트리에 대하여 알아보고 관련 문제를 더 풀어보기 위해 이 문제를 풀어보게 되었는데... 대충 a를 기준으로 겹쳐지는 선의 수를 세어나가면 되겠다는 생각을 했는데 세그먼트 트리를 어떻게 그려나가야할지를 잘 모르겠었다.. 그래서 검색을 통하여 문제를 어떻게 접근하는지를 보았다.

대략적인 순서는 다음과 같다.

1. A를 배열에 넣어준다.
2. B는 딕셔너리에 넣어주는데 key를 식별번호로, value를 index로 저장을 해준다.
3. 세그먼트 트리로 사용하기 위한 tree를 선언해준다.
4. A배열을 기준으로 처음부터 끝까지 for문으로 돌려준다.
5. B[A의 식별번호]의 인덱스를 방문했을 때 겹치는 선의 수를 세준다.
6. B[A의 식별번호]의 인덱스를 방문했다는 의미로 세그먼트 트리를 update해준다.

그림으로 선을 하나씩 추가할 때마다의 상황을 그려보았다.

겹치는 선들을 세어나가는 과정은 위와 같고, 아래는 세그먼트 트리를 update를 해가면서 구간의 합들을 구해 저장해놓는 과정이다.

이 문제를 세그먼트 트리를 사용하지 않고 구하려면 방문을 했을때 1로 최신화하기 전에 그 index부터 끝까지 확인하며 모두 더하면 된다. 그러면 for문으로 그 index부터 끝까지 모두 탐색을 해서 더해주려면 시간복잡도가 너무 크게된다. 따라서 구간의 합을 구하는데 시간복잡도가 O(logN)인 세그먼트 트리를 사용해준다.

그리고 B를 저장할때 배열이 아니고 딕셔너리를 사용했는데, 이는 배열로 사용하면 A와 같은 식별번호의 인덱스값을 처음부터 탐색을 해서 찾아내야하지만 딕셔너리를 사용하면 식별번호를 키값으로 바로 인덱스값이 value로 나오기때문에 딕셔너리를 사용하였다.

코드

# 7578번 공장
from math import *
import sys

# 선을 그려 방문 업데이트
def update(node, start, end, idx):
    if idx < start or end < idx:
        return 0

    if start == end:
        tree[node] = 1
        return tree[node]

    mid = (start + end) // 2
    update(node*2,start,mid,idx)
    update(node*2+1,mid+1,end,idx)
    tree[node] = tree[node*2] + tree[node*2+1]
    return tree[node]


# 구간합 쿼리
def query(node, start, end, left, right):
    if right < start or end < left:
        return 0

    if left <= start and end <= right:
        return tree[node]

    mid = (start + end) // 2
    return query(node*2,start,mid,left,right) + query(node*2+1,mid+1,end,left,right)

# main
n = int(sys.stdin.readline())
arrA = sys.stdin.readline().split()

# arrB dict 초기화
idx = 0
arrB = {}
for num in sys.stdin.readline().split():
   arrB[num] = idx
   idx += 1

# 세그먼트 트리
h = int(ceil(log2(n)))
t_size = 1 << (h+1)

tree = [0] * t_size

answer = 0
for num in arrA:
    s_idx = arrB[num]
    answer += query(1,0,n-1,s_idx,n-1)

    update(1,0,n-1,s_idx)

print(answer)