[백준] 다도해 19593 python 풀이

https://www.acmicpc.net/problem/19593

핵심

• 어느시점에서 루프를 끝내야 하는지
  • 모든 섬이 연결되었을때
  • 또는 더이상 돌아도 모든섬을 연결할 수 없을때 ⇒ 어떻게 판단?
  • E[ i ] = (E[ i-1 ] * A + B) % N^2 ⇒ 점화식이며 모듈러 N^2 이므로 N제곱이내의 수
  • 다시말해 기존 방문한 곳을 다시방문하는 순간 루프 break
• 모든 섬이 연결되었는지를 어떻게 판단
  • 새로운 섬을 연결할 때마다 탐색하면 시간초과
  • 유니온 파인드를 통해 판단
  • 문제는 예를들어 4개의 섬(0, 1, 2 ,3)이 있을때
    • 0, 1 연결 ⇒ parents list : [0, 0, 2, 3]
    • 2 ,3 연결 ⇒ parents list : [0, 0, 2, 2]
    • 1, 3 연결 ⇒ parents list : [0, 0, 2, 0]
      • 이순간에 전체 연결이 이뤄짐 (2번섬도 연결이 이뤄짐. 3번섬을 통해 0번섬과 이어짐)
      • union 함수에서 매번 두 섬을 연결할때마다 리스트를 전체 돌면서(N번) 2인 값을 가진 모든 값을 0으로 바꿔줌
      • 너무 느림 ~ 더 나은 방법이 없을까?

    
    """
    Title : 다도해
    Link : https://www.acmicpc.net/problem/19593
    """
    import sys
    input = sys.stdin.readline
    
    def find_parent(a, parents):
        if a == parents[a]:
            return a
        parents[a] = find_parent(parents[a], parents)
        return parents[a]
    
    def union(a, b, parents):
        pa = find_parent(a, parents)
        pb = find_parent(b, parents)
        if pa == pb:
            return
        if pa < pb:
            for i in range(N):
                if parents[i] == pb:
                    parents[i] = pa
        else:
            for i in range(N):
                if parents[i] == pa:
                    parents[i] = pb
        return
    
    def solve(N, Seed, A, B):
        X, Y = 0, 0
        Nsq = N * N
        parents = [i for i in range(N)]
        isVisited = [False for _ in range(Nsq)]
        answer = 0
    
        E = Seed % Nsq
    
        while True:
            if isVisited[E] == True:
                return 0
            isVisited[E] = True
            answer += 1
            X = E // N
            Y = E % N
            if X == Y:
                E = (E * A + B) % Nsq
                continue
            # Union Find
            union(X, Y, parents)
            # check if all islands are connected
            cnt = 0
            for i in range(N):
                if parents[i] == 0:
                    cnt += 1
            if cnt == N:
                break
            E = (E * A + B) % Nsq
        return answer
    
    t = int(input())
    for _ in range(t):
        N, Seed, A, B = map(int, input().split())
        print(solve(N, Seed, A, B))

[추가수정] 속도 개선방법

사이클이 생기지 않게 노드를 연결하면 전체 노드개수(N)의 N-1개를 연결하면 반드시 전체 노드가 연결이 됨.
따라서 union함수에서 연결이 이뤄줬는지 리턴값으로 주고
connections 변수로 총연결개수를 계산해서 N-1 이상이면 브레이크 하도록 해야함.

import sys
input = sys.stdin.readline


def find_parent(a, parents):
    if a == parents[a]:
        return a
    parents[a] = find_parent(parents[a], parents)
    return parents[a]

def union(a, b, parents):
    pa = find_parent(a, parents)
    pb = find_parent(b, parents)
    if pa == pb:
        return 0
    if pa < pb:
        parents[pb] = pa
    else:
        parents[pa] = pb
    return 1


def solve(N, Seed, A, B):
    X, Y = 0, 0
    Nsq = N * N
    parents = [i for i in range(N)]
    isVisited = [False for _ in range(Nsq)]
    answer = 0

    E = Seed % Nsq
    connections = 0
    while True:
        if isVisited[E] == True:
            return 0
        isVisited[E] = True
        answer += 1
        X = E // N
        Y = E % N
        if X == Y:
            E = (E * A + B) % Nsq
            continue
        # Union Find
        connections += union(X, Y, parents)
        # check all islands are connected
        if connections >= N-1:
            break
        E = (E * A + B) % Nsq
    return answer


t = int(input())
for _ in range(t):
    N, Seed, A, B = map(int, input().split())
    print(solve(N, Seed, A, B))