백준 #25187 고인물이 싫어요

코드

import sys
# sys.stdin = open("input.txt")
input = sys.stdin.readline


'''
n: 물탱크의 수 
m: 파이프의 수
q: 물탱크에 방문할 횟수
'''

def find(x):
    if x != parent[x]:
        return find(parent[x])
    return parent[x]


def union(x, y):
    x = find(x)
    y = find(y)

    if x < y:
        parent[y] = x
        water[x] += water[y]
    else:
        parent[x] = y
        water[y] += water[x]

n, m, q = map(int, input().split())
arr = [0] + list(map(int, input().rstrip().split()))    # 청정수: 1 고인물: 0
water = [0] * (n + 1)
for i in range(1, n + 1):
    if arr[i]:
        water[i] = 1
    else:
        water[i] = -1
parent = [i for i in range(n + 1)]
for _ in range(m):
    a, b = map(int, input().split())
    if find(a) != find(b):
        union(a, b)


for _ in range(q):
    v = int(input())    # 방문할 물탱크 번호
    if water[find(v)] > 0:
        print(1)
    else:
        print(0)

틀린 코드

import sys
# sys.stdin = open("input.txt")
input = sys.stdin.readline


'''
n: 물탱크의 수 
m: 파이프의 수
q: 물탱크에 방문할 횟수
'''

def find(x):
    if x != parent[x]:
        return find(parent[x])
    return parent[x]


def union(x, y):
    x = find(x)
    y = find(y)

    if x < y:
        parent[y] = x
    else:
        parent[x] = y


n, m, q = map(int, input().split())
water = [0] + list(map(int, input().rstrip().split()))    # 청정수: 1 고인물: 0
parent = [i for i in range(n + 1)]
for _ in range(m):
    a, b = map(int, input().split())
    if parent[a] != parent[b]:
        union(a, b)

visited = [-1] * (n + 1)
for _ in range(q):
    v = int(input())    # 방문할 물탱크 번호
    parent_v = parent[v]
    if visited[parent_v] == -1:
        good = 0
        bad = 0
        for i in range(1, n + 1):
            if parent_v == parent[i]:
                if water[i]:
                    good += 1
                else:
                    bad += 1
        if good > bad:
            print(1)
            visited[parent_v] = 1
        else:
            print(0)
            visited[parent_v] = 0
    else:
        print(visited[parent_v])

풀이

금방 풀릴 줄 알았는데 시간초과 때문에 애를 많이 먹었다 ㅠㅠ

N과 Q 범위가 각각 100,000까지 가능하기 때문에 틀린 코드에서 보면 당연히 시간 초과가 날 수 밖에 없다…

그래도 visited 배열을 써서 어떻게든 줄여보려고 했는데 계속 실패했다.

몇 달 전에도 유니온 파인드 문제를 풀다가 시간초과때문에 애를 먹었던 기억이 떠올라서 그 문제를 찾아가 보았다.

백준 4195번 친구 네트워크 문제인데 friend라는 배열에 친구 네트워크의 수를 저장해놓았다가 union단계에서 합쳐주는 방식으로 풀었었다.

이 문제에서는 water이라는 배열을 만들어서 초기값을 청정수이면 1을 고인물이면 -1로 설정해준 다음에 union단계에서 부모 물탱크에 자식 물탱크 값을 더해주었다.

또한 중복되는 물탱크 쌍이 주어질 수 있기때문에 find 값이 다를 경우에만 union을 수행해 주었다.

문제 난이도가 올라갈 수록 시간이나 메모리에 대한 신경을 많이 써주어야하는데 아직 이부분에 대해 많이 부족한 것 같다 ㅠㅠ

더 열심히 하자!! 화이팅!!!