백준 10838번. 트리(python)

문제 : https://www.acmicpc.net/problem/10838

풀이

LCA문제인데 depth를 사용하면 안된다.

또한, paint와 count 연산 시 a번 노드와 b번 노드 사이의 최단경로의 길이는 항상 1,000 이하이다.

위 문구를 통해 문제를 해결해야 한다. 즉, LCA를 반복문으로 진행해서 구해내야 한다.
부모노드와 간선의 색상(부모노드a에서 자식노드b로 이어지는 간선을 color[b]로 둔다)만 구해서 사용한다.

시간초과된 풀이(36%)

import sys
input=sys.stdin.readline

N,K=map(int,input().split())
son=[set() for _ in range(N)]
son[0]={i for i in range(1,N)}
depth=[1]*N
parent=[0]*N
color=[0]*N

depth[0]=0

def count_color(a,b):
    colors=set()
    while depth[a]!=depth[b]:
        if depth[a]>depth[b]:
            colors.add(color[a])
            a=parent[a]
        else:
            colors.add(color[b])
            b=parent[b]
    while a!=b:
        colors.add(color[a])
        colors.add(color[b])
        a=parent[a]
        b=parent[b]
    return len(colors)
def fix_depth(x,val):
    depth[x]=val
    if not son[x]:
        return
    for i in son[x]:
        fix_depth(i,val+1)

def fix_color(a,b,val):
    while depth[a]!=depth[b]:
        if depth[a]>depth[b]:
            color[a]=val
            a=parent[a]
        else:
            color[b]=val
            b=parent[b]
    while a!=b:
        color[a]=val
        color[b]=val
        a=parent[a]
        b=parent[b]

def lca(a,b):
    while depth[a]!=depth[b]:
        if depth[a]>depth[b]:
            a=parent[a]
        else:
            b=parent[b]
    while a!=b:
        a=parent[a]
        b=parent[b]
    return a

for _ in range(K):
    tmp=list(map(int,input().split()))
    r=tmp[0]
    if r==1: #칠하기
        a,b,c=tmp[1:]
        fix_color(a,b,c)
    elif r==2: # 노드변경
        a,b=tmp[1:]
        son[parent[a]].discard(a)
        parent[a]=b
        son[b].add(a)
        fix_depth(b,depth[b])
    else: # 사이의 다른 색상 구하기
        a,b=tmp[1:]
        print(count_color(a,b))

정답코드

import sys
input=sys.stdin.readline

N,K=map(int,input().split())

parent=[0]*N
color=[0]*N

def find_lca(a,b):
    tmp=set()
    for i in range(1000):
        tmp.add(a)
        a=parent[a]
    for i in range(1000):
        if b in tmp:
            return b
        b=parent[b]
        
def count_color(a,b):
    colors=set()
    top=find_lca(a,b)
    while top!=a:
        colors.add(color[a])
        a=parent[a]
    while top!=b:
        colors.add(color[b])
        b=parent[b]
    return len(colors)

def fix_color(a,b,val):
    top=find_lca(a,b)
    while top!=a:
        color[a]=val
        a=parent[a]
    while top!=b:
        color[b]=val
        b=parent[b]

for _ in range(K):
    tmp=list(map(int,input().split()))
    r=tmp[0]
    if r==1: #칠하기
        a,b,c=tmp[1:]
        fix_color(a,b,c)
    elif r==2: # 노드변경
        a,b=tmp[1:]
        parent[a]=b
    else: # 사이의 다른 색상 구하기
        a,b=tmp[1:]
        print(count_color(a,b))