Notice
Recent Posts
Recent Comments
Link
«   2026/05   »
1 2
3 4 5 6 7 8 9
10 11 12 13 14 15 16
17 18 19 20 21 22 23
24 25 26 27 28 29 30
31
Archives
Today
Total
관리 메뉴

mo1lusca의 블로그

[Unifox] 동아리 프로젝트 - 트리에서의 구간 쿼리 알고리즘 시각화 본문

Unifox

[Unifox] 동아리 프로젝트 - 트리에서의 구간 쿼리 알고리즘 시각화

mo1lusca 2025. 12. 7. 16:12

목차 

  1. 주제 선정 동기
  2. 개념 설명
    • 선형 구조에서의 구간 쿼리
    • 비선형 구조에서의 구간 쿼리
    • ETT, HLD
  3. 주요 코드 설명
  4. 배운점, 느낀점 및 개선점

1. 주제 선정 동기

동아리에서 처음 트리 자료구조에 대해 배우고 흥미가 생겨, 여름방학과 2학기 중 개인적으로 트리에 대해 공부했다.

트리에서의 구간 쿼리와 관련한 알고리즘도 공부했는데, 아이디어가 상당히 재밌고 시각화하면

이해하기도 훨씬 쉬울 것 같아 주제로 선정하게 되었다.


2. 개념 설명

i. 선형 구조에서의 구간 쿼리

구간 쿼리란, 자료구조 내의 연속된 범위에 대해 업데이트/질의(쿼리)를 수행하는 것이다.

쉽게 설명해 보겠다.. { 1, 3, 2, 4, 5 } 라는 배열이 있을 때, [1,3] 구간의 합을 구하는 쿼리를 수행하면, 3+2+4=9라는 결과를 반환하고... 대충 이런 것이다.

방금 말했듯 구간 쿼리는 당연하게도 구간에 대해 쿼리를 수행하기 때문에, 연속된 범위로 표현되는 구간이 존재해야 한다.

이러한 성질 때문에 대부분의 구간 쿼리는 선형구조 (배열, 리스트 등)에서 수행된다.

선형 구조에서 구간 쿼리를 수행하는 대표적인 알고리즘으로는.. 누적합, 세그먼트 트리, Mo's 등이 있다.

 

여기서 세그먼트 트리라는 자료구조를 기억해두어야 한다. 세그먼트 트리란 이진트리를 기반으로 만들어진 자료구조이다.

선형 구조(배열)의 각 원소들을 리프 노드로 가지며, 내부 노드는 자신의 두 자식의 값을 연산한 결과를 값으로 가진다.

예를 들어 1, 4라는 값을 가진 연속된 두 리프 노드가 있고, 구간합을 구하려고 한다면, 이 두 리프 노드를 동시에 자식으로 가지는 부모 노드에는 5(=1+4)라는 값이 들어간다.

 

ii. 비선형 구조에서의 구간 쿼리

그렇다면 비선형 구조에서는 어떻게 구간 쿼리를 수행할 수 있을까?

일단 간단히 생각해 보면 그래프에서는 사이클이 있기 때문에 구간을 잘 정의 할 수 없다..

그러므로 여기서 말하는 비선형 구조는 결국 트리라고 생각해도 된다.

다시 본론으로 돌아와서, 구간 쿼리를 수행하기 전에 트리에서의 구간이 무엇인가 알 필요가 있다.

조금 생각해보면 트리에서의 구간의 기준은 크게 두 가지로 나눠짐을 알 수 있다.

첫 번째는 임의의 노드 u, v에 대한 경로이다. 트리에서는 임의의 한 정점에서 다른 한 정점까지의 경로가 유일하므로,

(u, v) 경로는 구간이라고 생각할 수 있다.

두 번째는 임의의 노드 u를 루트 노드로 삼는 서브트리이다. 

 

그래서 이런 기준을 가지고 구간 쿼리를 어떻게 수행하라는 것인가?

매우 직관적으로 생각해서, 트리를 선형구조로 쫙 펴버리면 된다.

이때, 트리에서의 경로를 구간으로 나타내는 알고리즘이 HLD, 서브트리를 구간으로 나타내는 알고리즘이 ETT이다.

 

iii. ETT, HLD

ETT란, Euler Tour Technique의 약자로, DFS Ordering이라고도 부른다.

ETT는 트리에서의 임의의 서브트리를 연속된 구간으로 나타낼 수 있게 해 준다.

구체적으로는, DFS를 수행하며 각 노드별로 in, out값을 가지게 한다.

in은 DFS를 수행하며 해당 노드를 처음으로 방문한 시점이고, out은 해당 노드의 모든 자손 노드가 탐색이 완료된 시점이다.

이렇게 in, out값을 잡으면 임의의 정점 u에 대해, u를 루트로 하는 서브트리의 모든 정점은 [ in[u], out[u] ] 내의 구간으로 표현된다.

 

HLD란, Heavy-Light Decomposition의 약자이다.

기본적인 아이디어는, 비선형 구조(트리)를 여러 개의 선형 구조로 분해하는 것이다.

이때 서브트리의 크기를 기준으로 간선을 나눈다.

정확히는, 부모 노드 u에서 자식 노드 v로 가는 간선 (u, v)가 있을 때, v의 서브트리의 크기가 u의 서브트리의 크기의 절반 이상인 경우 간선 (u, v)를 heavy edge라고 하고 노드 v를 heavy child 내지는 heavy node라고 부른다.

또한 heavy edge가 아닌 u와 연결된 나머지 간선은 light edge라고 부른다.

이런 식으로 heavy edge를 정의하면, 임의의 노드에서 자식으로 향하는 heavy edge는 최대 한 개 있게 된다.

또한 트리 내의 모든 노드가 heavy chain들 중 하나에는 속한다는 것을 알 수 있다. (자기 자신이 체인인 경우도 포함한다.)

 

그렇다면 임의의 노드에서 자식으로 향하는 heavy edge를 계속 따라간다고 해보자. 앞서 말한 특성 때문에 heavy edge로 이루어진 하나의 체인(chain)이 나오게 되고, 이를 heavy chain이라 부른다.

heavy chain은 선형구조이고, 트리 내의 모든 경로는 각기 다른 부분 heavy chain의 집합으로 표현할 수 있으므로

트리 내의 모든 경로에 대해 선형구조처럼 구간 쿼리를 날릴 수 있다.

쉽게 말하자면, 트리를 적당히 잘라서 선형으로 만든 다음 연산한다는 것이다.


3. 주요 코드 설명

tree.py

세그먼트 트리, HLD, ETT를 관리한다.

 

더보기
class SegTree:
    def __init__(self, n):
        self.n = n
        self.tree = [0] * (4 * n + 16)
        self.lazy = [0] * (4 * n + 16)
    
    def push(self, idx, left, right):
        if self.lazy[idx] != 0:
            self.tree[idx] += self.lazy[idx] * (right - left + 1)
            if left != right:
                self.lazy[idx*2] += self.lazy[idx]
                self.lazy[idx*2 + 1] += self.lazy[idx]
            self.lazy[idx] = 0

    def update(self, idx, left, right, cleft, cright, cval):
        self.push(idx, left, right)
        if cright < left or right < cleft:
            return
        if cleft <= left and right <= cright:
            self.lazy[idx] += cval
            self.push(idx, left, right)
            return
        mid = (left + right) // 2
        self.update(idx*2, left, mid, cleft, cright, cval)
        self.update(idx*2+1, mid+1, right, cleft, cright, cval)
        self.push(idx*2, left, mid)
        self.push(idx*2+1, mid+1, right)
        self.tree[idx] = self.tree[idx*2] + self.tree[idx*2+1]
        
    def query(self, idx, left, right, qleft, qright):
        self.push(idx, left, right)
        if qright < left or right < qleft:
            return 0
        if qleft <= left and right <= qright:
            return self.tree[idx]
        mid = (left + right) // 2
        return self.query(idx*2, left, mid, qleft, qright) + self.query(idx*2+1, mid+1, right, qleft, qright)

구간 업데이트 & 구간 쿼리 합연산 세그트리를 구현했다.

명색이 시각화(시뮬레이터)인데 점 업데이트면 너무 없어보여서... lazy하게 갱신해주자

 

더보기
class HLD:
    # HLD용 클래스
    def __init__(self, n):
        # HLD 구현용
        self.n = n
        self.size = [0] * (n + 1)        # 서브트리 크기
        self.dep = [0] * (n + 1)         # 깊이
        self.par = [0] * (n + 1)         # 부모 노드
        self.top = [0] * (n + 1)         # 체인의 최상단 정점
        self.in_time = [0] * (n + 1)     # ETT 시작 시간
        self.out_time = [0] * (n + 1)    # ETT 종료 시간
        self.g = [[] for _ in range(n + 1)]  # 그래프
        self.cnt = 0                     # ETT용
        self.seg = SegTree(n)
        
        # 시각화용
        self.id = [0] * (n + 1)          # 체인 ID
        self.chain_cnt = 0               # 체인 개수
        self.val = [0] * (n + 1)         # 노드 값
    
    def add_edge(self, u, v):
        # 간선 추가
        self.g[u].append(v)
        self.g[v].append(u)
    
    def set_value(self, v, value):
        # main에서 이 함수를 통해 값을 넘겨줄 것임
        self.val[v] = value
    
    def dfs1(self, v=1, p=0):
        # par, size, dep 저장, heavy 간선 앞으로 보내기
        self.par[v] = p
        self.size[v] = 1
        for i in range(len(self.g[v])):
            child = self.g[v][i]
            if child == p:
                continue
            self.dep[child] = self.dep[v] + 1
            self.par[child] = v
            self.dfs1(child, v)
            self.size[v] += self.size[child]
            if self.g[v][0] == p or self.size[child] > self.size[self.g[v][0]]:
                self.g[v][0], self.g[v][i] = self.g[v][i], self.g[v][0]
    
    def dfs2(self, v=1, p=0):
        # dfs ordering
        self.cnt += 1
        self.in_time[v] = self.cnt
        
        # 체인 ID 부여
        if v == self.top[v]:
            self.chain_cnt += 1
        self.id[v] = self.chain_cnt
        
        for child in self.g[v]:
            if child == p:
                continue
            # Heavy edge면 같은 체인, 아니면 새로운 체인 시작
            self.top[child] = self.top[v] if child == self.g[v][0] else child
            self.dfs2(child, v)
        self.out_time[v] = self.cnt
    
    def init(self):
        # HLD 초기화
        self.dep[1] = 0
        self.par[1] = 0
        self.top[1] = 1
        self.dfs1()
        self.dfs2()
        
        # 입력받은 값으로 세그트리 리프노드 채우기
        for v in range(1, self.n + 1):
            self.seg.update(1,1,self.n,self.in_time[v],self.in_time[v], self.val[v])

HLD 사전 설정? 함수들이다. 

일반적인 HLD 구현 코드와 딱히 다를점은 없고, set_value 함수를 통해 main.py에서 입력을 받는 정도만 기억하자.

 

더보기
    def update(self, u, v, val):
        affected_nodes = []
        
        while self.top[u] != self.top[v]:
            if self.dep[self.top[u]] < self.dep[self.top[v]]:
                u, v = v, u
            head = self.top[u]
            
            for node in range(1,self.n+1):
                if self.in_time[head] <= self.in_time[node] <= self.in_time[u]:
                    affected_nodes += [node]
                    self.val[node] += val
                
            yield (affected_nodes.copy(), 'UPDATE', val)
            self.seg.update(1, 1, self.n, self.in_time[head], self.in_time[u], val)
            u = self.par[head]
        
        if self.dep[u] > self.dep[v]:
            u, v = v, u
        
        for node in range(1,self.n+1):
            if self.in_time[u] <= self.in_time[node] <= self.in_time[v]:
                affected_nodes += [node]
                self.val[node] += val

        yield (affected_nodes.copy(), 'UPDATE', val)
        self.seg.update(1, 1, self.n, self.in_time[u], self.in_time[v], val)
    
    def query(self, u, v):
        affected_nodes = []
        res = 0
        
        while self.top[u] != self.top[v]:
            if self.dep[self.top[u]] < self.dep[self.top[v]]:
                u, v = v, u
            head = self.top[u]
            
            for node in range(1,self.n+1):
                if self.in_time[head] <= self.in_time[node] <= self.in_time[u]:
                    affected_nodes += [node]
            
            res += self.seg.query(1, 1, self.n, self.in_time[head], self.in_time[u])
            yield (affected_nodes.copy(), 'QUERY', res)
            u = self.par[head]
        
        if self.dep[u] > self.dep[v]:
            u, v = v, u
        
        for node in range(1,self.n+1):
            if self.in_time[u] <= self.in_time[node] <= self.in_time[v]:
                affected_nodes += [node]

        res += self.seg.query(1, 1, self.n, self.in_time[u], self.in_time[v])
        yield (affected_nodes.copy(), 'QUERY', res)
    
    def is_heavy(self, u, v):
        # heavy child인지 검사
        if self.dep[u] > self.dep[v]:
            u, v = v, u
        return self.par[v] == u and len(self.g[u]) > 0 and self.g[u][0] == v

chain들을 엮어 경로 업데이트 / 쿼리를 진행하는 함수들이다. 전체적인 구조는 일반적인 HLD 구현과 같지만,

yield를 통해서 영향을 받는 노드를 실시간으로 반환한다. (언어를 파이썬으로 선택한 이유이기도 하다)

대략적인 흐름을 설명하자면,

경로를 구하려는 두 노드가 속한 chain중 더 깊은 위치에 있는 chain 선택

--> chain의 head(최상단 노드)까지 구간 쿼리

--> 해당 head의 부모가 속한 chain과, 아직 건드리지 않은 chain을 대상으로 앞의 절차 반복

------> 하나의 chain만 남게 된다면 마지막으로 구간 쿼리 후 종료

 

더보기
class ETT:
    # ETT용 클래스
    def __init__(self, n):
        self.n = n
        self.par = [0] * (n + 1)
        self.dep = [0] * (n + 1)
        self.in_time = [0] * (n + 1)
        self.out_time = [0] * (n + 1)
        self.g = [[] for _ in range(n + 1)]
        self.cnt = 0
        self.seg = SegTree(n)
        self.val = [0] * (n + 1)         # 노드 값
    
    def add_edge(self, u, v):
        self.g[u].append(v)
        self.g[v].append(u)
    
    def set_value(self, v, value):
        self.val[v] = value
    
    def dfs(self, v=1, p=0):
        # dfs ordering
        self.par[v] = p

        self.cnt += 1
        self.in_time[v] = self.cnt

        for child in self.g[v]:
            if child == p:
                continue

            self.dep[child] = self.dep[v] + 1
            self.dfs(child, v)

        self.out_time[v] = self.cnt
    
    def init(self):
        # ETT 초기화
        self.dep[1] = 0
        self.par[1] = 0
        self.dfs()
        
        for v in range(1, self.n + 1):
            self.seg.update(1,1,self.n, self.in_time[v],self.in_time[v],self.val[v])

ETT 사전 설정 부분이다.

dfs 함수가 핵심으로, 2. iii에서 설명한것과 동일하게 in, out 배열을 채운다.

 

더보기
    def query(self, v):
        affected_nodes = []
        
        # v를 루트로 하는 서브트리의 모든 노드 찾기
        def collect_subtree(node):
            affected_nodes.append(node)
            for child in self.g[node]:
                if child != self.par[node]:
                    collect_subtree(child)
        
        collect_subtree(v)
        result = self.seg.query(1, 1, self.n, self.in_time[v], self.out_time[v])
        yield (affected_nodes, 'QUERY', result)
    
    def update(self, v, val):
        affected_nodes = []
        
        # v를 루트로 하는 서브트리의 모든 노드 찾기
        def collect_subtree(node):
            affected_nodes.append(node)
            self.val[node] += val
            for child in self.g[node]:
                if child != self.par[node]:
                    collect_subtree(child)
        
        collect_subtree(v)
        self.seg.update(1, 1, self.n, self.in_time[v], self.out_time[v], val)
        yield (affected_nodes, 'UPDATE', val)
    
    def is_ancestor(self, u, v):
        #u가 v의 조상인지 확인
        return self.in_time[u] <= self.in_time[v] <= self.out_time[u]

서브트리를 구간으로 업데이트 / 쿼리를 진행하는 함수이다.

상당히 간단한 구조로 동작하는데, 위에서 말했듯 ETT는 HLD와 다르게 서브트리를 바로 구간으로

표현할 수 있기 때문에, 쿼리를 한번만 날려도 되기 때문이다.

내부함수 collect_subtree는 시각화를 위해 쿼리에 영향을 받는 노드를 찾는 용도이다.

 


render.py

트리와, 쿼리 과정을 화면에 띄워준다.

 

더보기
    def draw_edge(self, u, v, color, width, highlight=False):
        # 간선 그리기
        if u not in self.positions or v not in self.positions:
            return
        
        pos_u = self.positions[u]
        pos_v = self.positions[v]
        
        if highlight:
            # 하이라이트 효과
            highlight_color = HIGHLIGHT_COLOR_QUERY if self.highlight_type == 'QUERY' else HIGHLIGHT_COLOR_UPDATE
            pygame.draw.line(self.screen, highlight_color, pos_u, pos_v, width + 6)
        
        pygame.draw.line(self.screen, color, pos_u, pos_v, width)
    
    def draw_node(self, v, color, highlight=False):
        # 노드 그리기
        if v not in self.positions:
            return
        
        pos = self.positions[v]
        
        # 하이라이트 효과
        if highlight:
            highlight_color = HIGHLIGHT_COLOR_QUERY if self.highlight_type == 'QUERY' else HIGHLIGHT_COLOR_UPDATE
            pygame.draw.circle(self.screen, highlight_color, pos, NODE_RADIUS + 7)
        
        # 노드 외곽선
        pygame.draw.circle(self.screen, BLACK, pos, NODE_RADIUS + 2)
        # 노드 내부
        pygame.draw.circle(self.screen, color, pos, NODE_RADIUS)
        
        # 노드 값 표시
        value_text = self.value_font.render(str(self.tree.val[v]), True, WHITE)
        value_rect = value_text.get_rect(center=pos)
        self.screen.blit(value_text, value_rect)
        
        # 노드 번호는 작게 위에 표시
        label_text = self.font.render(f"#{v}", True, BLACK)
        label_pos = (pos[0] - 15, pos[1] - NODE_RADIUS - 20)
        self.screen.blit(label_text, label_pos)
    
    def draw_tree(self):
        # 전체 트리 그리기

        # 간선 먼저 그리기
        for v in range(1, self.tree.n + 1):
            parent = self.tree.par[v]
            if parent != 0:
                edge_color = LIGHT_COLOR
                edge_width = EDGE_WIDTH_LIGHT
                
                # HLD의 경우 heavy edge 강조
                if self.is_hld and self.tree.is_heavy(parent, v):
                    chain_id = self.tree.id[v]
                    color_idx = (chain_id - 1) % len(CHAIN_COLORS)
                    edge_color = CHAIN_COLORS[color_idx]
                    edge_width = EDGE_WIDTH_HEAVY
                
                # 간선 하이라이트: 양쪽 노드가 모두 하이라이트되어 있을 때만
                highlight = v in self.highlight_nodes and parent in self.highlight_nodes
                self.draw_edge(parent, v, edge_color, edge_width, highlight)
        
        # 노드 그리기
        for v in range(1, self.tree.n + 1):
            node_color = BLUE
            
            # HLD의 경우 체인별 색상
            if self.is_hld:
                chain_id = self.tree.id[v]
                color_idx = (chain_id - 1) % len(CHAIN_COLORS)
                node_color = CHAIN_COLORS[color_idx]
            
            highlight = v in self.highlight_nodes
            self.draw_node(v, node_color, highlight)

간선, 노드, 트리를 그리는 함수이다. hld인 경우 노드색을 chainID별로 다르게 했고,

쿼리에 영향을 받는 노드 / 간선인 경우 쿼리 종류에 따라 다르게 하이라이트 처리를 해줬다.

 

더보기
    def run(self):
        running = True
        selected_node = 1
        clock = pygame.time.Clock()
        
        # 쿼리 모드 상태
        query_mode = None  # None, 'UPDATE', 'QUERY'
        query_state = 'IDLE'  # 'IDLE', 'SELECT_FIRST', 'SELECT_SECOND', 'INPUT_VALUE', 'EXECUTING'
        query_node1 = 0
        query_node2 = 0
        query_value = ""
        query_generator = None
        
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        if query_state != 'IDLE':
                            # 쿼리 모드 취소
                            query_mode = None
                            query_state = 'IDLE'
                            query_node1 = 0
                            query_node2 = 0
                            query_value = ""
                            self.set_highlight([])
                        else:
                            running = False
                    
                    # 쿼리 모드 시작
                    elif event.key == pygame.K_u and query_state == 'IDLE':
                        query_mode = 'UPDATE'
                        query_state = 'SELECT_FIRST'
                        query_value = ""
                    
                    elif event.key == pygame.K_q and query_state == 'IDLE':
                        query_mode = 'QUERY'
                        query_state = 'SELECT_FIRST'
                    
                    # 값 입력 모드
                    elif query_state == 'INPUT_VALUE':
                        if event.key == pygame.K_RETURN:
                            if query_value:
                                val = int(query_value)
                                if self.is_hld:
                                    if query_mode == 'UPDATE':
                                        query_generator = self.tree.update(query_node1, query_node2, val)
                                    else:
                                        query_generator = self.tree.query(query_node1, query_node2)
                                else:
                                    # ETT는 단일 노드만 사용
                                    if query_mode == 'UPDATE':
                                        query_generator = self.tree.update(query_node1, val)
                                    else:
                                        query_generator = self.tree.query(query_node1)
                                query_state = 'EXECUTING'
                                self.query_result = None
                        elif event.key == pygame.K_BACKSPACE:
                            query_value = query_value[:-1]
                        elif event.unicode.isdigit() or (event.unicode == '-' and not query_value):
                            query_value += event.unicode
                    # 실행중
                    elif query_state == 'EXECUTING' and event.key == pygame.K_RETURN:
                        if query_generator:
                            try:
                                nodes, op_type, result = next(query_generator)
                                self.set_highlight(nodes, op_type)
                                self.query_result = result
                            except StopIteration:
                                # 쿼리 완료
                                query_mode = None
                                query_state = 'IDLE'
                                query_generator = None
                                self.set_highlight([])
                                self.query_result = None
                
                # 노드 클릭
                if event.type == pygame.MOUSEBUTTONDOWN:
                    mouse_pos = pygame.mouse.get_pos()
                    clicked_node = self.find_clicked_node(mouse_pos)
                    
                    if clicked_node > 0:
                        if query_state == 'SELECT_FIRST':
                            query_node1 = clicked_node
                            # ETT는 단일 노드만 필요
                            if self.is_hld and (query_mode == 'QUERY' or query_mode == 'UPDATE'):
                                query_state = 'SELECT_SECOND'
                            else:
                                # ETT는 바로 값 입력 또는 쿼리 실행
                                if query_mode == 'UPDATE':
                                    query_state = 'INPUT_VALUE'
                                else:
                                    # 쿼리는 바로 실행
                                    query_generator = self.tree.query(query_node1)
                                    query_state = 'EXECUTING'
                            selected_node = clicked_node
                        elif query_state == 'SELECT_SECOND':
                            query_node2 = clicked_node
                            if query_mode == 'UPDATE':
                                query_state = 'INPUT_VALUE'
                            else:
                                # 쿼리는 바로 실행
                                if self.is_hld:
                                    query_generator = self.tree.query(query_node1, query_node2)
                                query_state = 'EXECUTING'
                            selected_node = clicked_node
                        else:
                            selected_node = clicked_node
            
            # 화면 그리기
            self.screen.fill(WHITE)
            self.draw_title()
            self.draw_tree()
            self.draw_info(selected_node)
            
            if query_state != 'IDLE':
                self.draw_query_panel(query_state, query_node1, query_node2, query_value)
            
            pygame.display.flip()
            clock.tick(60)
        
        pygame.quit()

전체적인 구동을 담당하는 함수이다.. 사실 코드 길이에 비해 엄청난 핵심 코드는 아니기 때문에 가볍게 보고 넘어가자..

 


main.py

콘솔에서 모드를 입력받고 파이게임 화면을 띄운다.

 

더보기
from tree import HLD, ETT
from render import TreeRenderer

def create_tree():
    n = 13
    edges = [
        (1, 2),
        (1, 3),
        (1, 4),
        (2, 5),
        (2, 6),
        (3, 7),
        (3, 8),
        (4, 9),
        (6, 10),
        (6, 11),
        (8, 12),
        (9, 13)
    ]
    return n, edges

def build_hld_tree(n, edges):
    hld = HLD(n)
    
    for v in range(1, n + 1):
        hld.set_value(v, v)
    
    for u, v in edges:
        hld.add_edge(u, v)
    
    hld.init()
    
    return hld


def build_ett_tree(n, edges):
    ett = ETT(n)

    for v in range(1, n + 1):
        ett.set_value(v, v)
    
    for u, v in edges:
        ett.add_edge(u, v)
    
    ett.init()
    
    return ett

객체나 트리를 사전설정해주는 함수들이다.. 어찌보면 당연하지만 트리는 하드코딩 되어있다.

 

더보기
def main():
    print("=" * 50)
    print("HLD & ETT 시각화 프로그램")
    print("=" * 50)
        
    n, edges = create_tree()

    # 시각화 모드 선택
    print("\n시각화 모드:")
    print("1. HLD (Heavy-Light Decomposition)")
    print("2. ETT (Euler Tour Tree)")
    
    mode = input("\n선택 (1-2, 기본값 1): ").strip()
    
    if mode == '2':
        # ETT 모드
        ett = build_ett_tree(n, edges)
        
        print("\n" + "=" * 50)
        print("ETT 시각화를 시작합니다...")
        print("=" * 50)
        
        renderer = TreeRenderer(ett, is_hld=False)
        renderer.run()
    else:
        # HLD 모드
        hld = build_hld_tree(n, edges)
        
        print("\n" + "=" * 50)
        print("HLD 시각화를 시작합니다...")
        print("=" * 50)
        
        renderer = TreeRenderer(hld, is_hld=True)
        renderer.run()

if __name__ == "__main__":
    main()

main함수이다. 어떤 모드를 사용할지 입력받고 모드에 따라 화면을 띄운다.

 


전체 코드

더보기
from tree import HLD, ETT
from render import TreeRenderer

def create_tree():
    n = 13
    edges = [
        (1, 2),
        (1, 3),
        (1, 4),
        (2, 5),
        (2, 6),
        (3, 7),
        (3, 8),
        (4, 9),
        (6, 10),
        (6, 11),
        (8, 12),
        (9, 13)
    ]
    return n, edges

def build_hld_tree(n, edges):
    hld = HLD(n)
    
    for v in range(1, n + 1):
        hld.set_value(v, v)
    
    for u, v in edges:
        hld.add_edge(u, v)
    
    hld.init()
    
    return hld


def build_ett_tree(n, edges):
    ett = ETT(n)

    for v in range(1, n + 1):
        ett.set_value(v, v)
    
    for u, v in edges:
        ett.add_edge(u, v)
    
    ett.init()
    
    return ett


def main():
    print("=" * 50)
    print("HLD & ETT 시각화 프로그램")
    print("=" * 50)
        
    n, edges = create_tree()

    # 시각화 모드 선택
    print("\n시각화 모드:")
    print("1. HLD (Heavy-Light Decomposition)")
    print("2. ETT (Euler Tour Tree)")
    
    mode = input("\n선택 (1-2, 기본값 1): ").strip()
    
    if mode == '2':
        # ETT 모드
        ett = build_ett_tree(n, edges)
        
        print("\n" + "=" * 50)
        print("ETT 시각화를 시작합니다...")
        print("=" * 50)
        
        renderer = TreeRenderer(ett, is_hld=False)
        renderer.run()
    else:
        # HLD 모드
        hld = build_hld_tree(n, edges)
        
        print("\n" + "=" * 50)
        print("HLD 시각화를 시작합니다...")
        print("=" * 50)
        
        renderer = TreeRenderer(hld, is_hld=True)
        renderer.run()


if __name__ == "__main__":
    main()

main.py

더보기
class SegTree:
    # 구간업데이트 구간쿼리 합 세그먼트 트리
    def __init__(self, n):
        self.n = n
        self.tree = [0] * (4 * n + 16)
        self.lazy = [0] * (4 * n + 16)
    
    def push(self, idx, left, right):
        if self.lazy[idx] != 0:
            self.tree[idx] += self.lazy[idx] * (right - left + 1)
            if left != right:
                self.lazy[idx*2] += self.lazy[idx]
                self.lazy[idx*2 + 1] += self.lazy[idx]
            self.lazy[idx] = 0

    def update(self, idx, left, right, cleft, cright, cval):
        self.push(idx, left, right)
        if cright < left or right < cleft:
            return
        if cleft <= left and right <= cright:
            self.lazy[idx] += cval
            self.push(idx, left, right)
            return
        mid = (left + right) // 2
        self.update(idx*2, left, mid, cleft, cright, cval)
        self.update(idx*2+1, mid+1, right, cleft, cright, cval)
        self.push(idx*2, left, mid)
        self.push(idx*2+1, mid+1, right)
        self.tree[idx] = self.tree[idx*2] + self.tree[idx*2+1]
        
    def query(self, idx, left, right, qleft, qright):
        self.push(idx, left, right)
        if qright < left or right < qleft:
            return 0
        if qleft <= left and right <= qright:
            return self.tree[idx]
        mid = (left + right) // 2
        return self.query(idx*2, left, mid, qleft, qright) + self.query(idx*2+1, mid+1, right, qleft, qright)


class HLD:
    # HLD용 클래스
    def __init__(self, n):
        # HLD 구현용
        self.n = n
        self.size = [0] * (n + 1)        # 서브트리 크기
        self.dep = [0] * (n + 1)         # 깊이
        self.par = [0] * (n + 1)         # 부모 노드
        self.top = [0] * (n + 1)         # 체인의 최상단 정점
        self.in_time = [0] * (n + 1)     # ETT 시작 시간
        self.out_time = [0] * (n + 1)    # ETT 종료 시간
        self.g = [[] for _ in range(n + 1)]  # 그래프
        self.cnt = 0                     # ETT용
        self.seg = SegTree(n)
        
        # 시각화용
        self.id = [0] * (n + 1)          # 체인 ID
        self.chain_cnt = 0               # 체인 개수
        self.val = [0] * (n + 1)         # 노드 값
    
    def add_edge(self, u, v):
        # 간선 추가
        self.g[u].append(v)
        self.g[v].append(u)
    
    def set_value(self, v, value):
        # main에서 이 함수를 통해 값을 넘겨줄 것임
        self.val[v] = value
    
    def dfs1(self, v=1, p=0):
        # par, size, dep 저장, heavy 간선 앞으로 보내기
        self.par[v] = p
        self.size[v] = 1
        for i in range(len(self.g[v])):
            child = self.g[v][i]
            if child == p:
                continue
            self.dep[child] = self.dep[v] + 1
            self.par[child] = v
            self.dfs1(child, v)
            self.size[v] += self.size[child]
            if self.g[v][0] == p or self.size[child] > self.size[self.g[v][0]]:
                self.g[v][0], self.g[v][i] = self.g[v][i], self.g[v][0]
    
    def dfs2(self, v=1, p=0):
        # dfs ordering
        self.cnt += 1
        self.in_time[v] = self.cnt
        
        # 체인 ID 부여
        if v == self.top[v]:
            self.chain_cnt += 1
        self.id[v] = self.chain_cnt
        
        for child in self.g[v]:
            if child == p:
                continue
            # Heavy edge면 같은 체인, 아니면 새로운 체인 시작
            self.top[child] = self.top[v] if child == self.g[v][0] else child
            self.dfs2(child, v)
        self.out_time[v] = self.cnt
    
    def init(self):
        # HLD 초기화
        self.dep[1] = 0
        self.par[1] = 0
        self.top[1] = 1
        self.dfs1()
        self.dfs2()
        
        # 입력받은 값으로 세그트리 리프노드 채우기
        for v in range(1, self.n + 1):
            self.seg.update(1,1,self.n,self.in_time[v],self.in_time[v], self.val[v])
        


    def update(self, u, v, val):
        affected_nodes = []
        
        while self.top[u] != self.top[v]:
            if self.dep[self.top[u]] < self.dep[self.top[v]]:
                u, v = v, u
            head = self.top[u]
            
            for node in range(1,self.n+1):
                if self.in_time[head] <= self.in_time[node] <= self.in_time[u]:
                    affected_nodes += [node]
                    self.val[node] += val
                
            yield (affected_nodes.copy(), 'UPDATE', val)
            self.seg.update(1, 1, self.n, self.in_time[head], self.in_time[u], val)
            u = self.par[head]
        
        if self.dep[u] > self.dep[v]:
            u, v = v, u
        

        for node in range(1,self.n+1):
            if self.in_time[u] <= self.in_time[node] <= self.in_time[v]:
                affected_nodes += [node]
                self.val[node] += val

        yield (affected_nodes.copy(), 'UPDATE', val)
        self.seg.update(1, 1, self.n, self.in_time[u], self.in_time[v], val)
    
    def query(self, u, v):
        affected_nodes = []
        res = 0
        
        while self.top[u] != self.top[v]:
            if self.dep[self.top[u]] < self.dep[self.top[v]]:
                u, v = v, u
            head = self.top[u]
            
            for node in range(1,self.n+1):
                if self.in_time[head] <= self.in_time[node] <= self.in_time[u]:
                    affected_nodes += [node]
            
            res += self.seg.query(1, 1, self.n, self.in_time[head], self.in_time[u])
            yield (affected_nodes.copy(), 'QUERY', res)
            u = self.par[head]
        
        if self.dep[u] > self.dep[v]:
            u, v = v, u
        
        for node in range(1,self.n+1):
            if self.in_time[u] <= self.in_time[node] <= self.in_time[v]:
                affected_nodes += [node]

        res += self.seg.query(1, 1, self.n, self.in_time[u], self.in_time[v])
        yield (affected_nodes.copy(), 'QUERY', res)
    

    def is_heavy(self, u, v):
        # heavy child인지 검사
        if self.dep[u] > self.dep[v]:
            u, v = v, u
        return self.par[v] == u and len(self.g[u]) > 0 and self.g[u][0] == v
    

class ETT:
    # ETT용 클래스
    def __init__(self, n):
        self.n = n
        self.par = [0] * (n + 1)
        self.dep = [0] * (n + 1)
        self.in_time = [0] * (n + 1)
        self.out_time = [0] * (n + 1)
        self.g = [[] for _ in range(n + 1)]
        self.cnt = 0
        self.seg = SegTree(n)
        self.val = [0] * (n + 1)         # 노드 값
    
    def add_edge(self, u, v):
        self.g[u].append(v)
        self.g[v].append(u)
    
    def set_value(self, v, value):
        # main에서 이 함수를 통해 값을 넘겨줌
        self.val[v] = value
    
    def dfs(self, v=1, p=0):
        # dfs ordering
        self.par[v] = p

        self.cnt += 1
        self.in_time[v] = self.cnt

        for child in self.g[v]:
            if child == p:
                continue

            self.dep[child] = self.dep[v] + 1
            self.dfs(child, v)

        self.out_time[v] = self.cnt
    
    def init(self):
        # ETT 초기화
        self.dep[1] = 0
        self.par[1] = 0
        self.dfs()
        
        for v in range(1, self.n + 1):
            self.seg.update(1,1,self.n, self.in_time[v],self.in_time[v],self.val[v])
    
    def query(self, v):
        affected_nodes = []
        
        # v를 루트로 하는 서브트리의 모든 노드 찾기
        def collect_subtree(node):
            affected_nodes.append(node)
            for child in self.g[node]:
                if child != self.par[node]:
                    collect_subtree(child)
        
        collect_subtree(v)
        result = self.seg.query(1, 1, self.n, self.in_time[v], self.out_time[v])
        yield (affected_nodes, 'QUERY', result)
    
    def update(self, v, val):
        affected_nodes = []
        
        # v를 루트로 하는 서브트리의 모든 노드 찾기
        def collect_subtree(node):
            affected_nodes.append(node)
            self.val[node] += val
            for child in self.g[node]:
                if child != self.par[node]:
                    collect_subtree(child)
        
        collect_subtree(v)
        self.seg.update(1, 1, self.n, self.in_time[v], self.out_time[v], val)
        yield (affected_nodes, 'UPDATE', val)
    
    def is_ancestor(self, u, v):
        #u가 v의 조상인지 확인
        return self.in_time[u] <= self.in_time[v] <= self.out_time[u]

tree.py

더보기
import pygame
from tree import HLD, ETT

# 색상 정의
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BLUE = (0, 0, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
YELLOW = (255, 255, 0)
ORANGE = (255, 165, 0)

# 체인별 색상
CHAIN_COLORS = [
    (255, 99, 71),    # Tomato
    (60, 179, 113),   # MediumSeaGreen
    (65, 105, 225),   # RoyalBlue
    (255, 165, 0),    # Orange
    (147, 112, 219),  # MediumPurple
    (0, 191, 255),    # DeepSkyBlue
    (255, 0, 255),    # Magenta
    (128, 0, 0),      # Maroon
    (255, 215, 0),    # Gold
    (46, 139, 87)     # SeaGreen
]

LIGHT_COLOR = (128, 128, 128)
HIGHLIGHT_COLOR_QUERY = (50, 205, 50)  # 파란색 - 쿼리
HIGHLIGHT_COLOR_UPDATE = (255, 140, 0)  # 주황색 - 업데이트

# 렌더링 설정
NODE_RADIUS = 25
EDGE_WIDTH_HEAVY = 4
EDGE_WIDTH_LIGHT = 2

SCREEN_WIDTH = 1920
SCREEN_HEIGHT = 1080


def init_pygame():
    #pygame 초기화
    pygame.init()
    # screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
    screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT), pygame.SCALED | pygame.FULLSCREEN)
    pygame.display.set_caption("HLD & ETT Visualization")
    return screen


class TreeRenderer:
    # 트리 렌더링해주는 클래스
    
    def __init__(self, tree_data, is_hld=True):
        self.tree = tree_data
        self.is_hld = is_hld
        self.screen = init_pygame()
        self.font = pygame.font.SysFont('malgungothic', 18)
        self.title_font = pygame.font.SysFont('malgungothic', 28, bold=True)
        self.value_font = pygame.font.SysFont('arial', 16, bold=True)
        
        self.positions = {}
        self.highlight_nodes = set()  # 하이라이트할 노드들
        self.highlight_type = 'QUERY'  # 'QUERY' or 'UPDATE'
        self.query_result = None  # 쿼리 결과
        
        self._calc_positions()
    
    def _calc_positions(self):
        if self.tree.n == 0:
            return
        
        # 각 깊이별 노드 개수 세기
        depth_nodes = {}
        for v in range(1, self.tree.n + 1):
            depth = self.tree.dep[v]
            if depth not in depth_nodes:
                depth_nodes[depth] = []
            depth_nodes[depth].append(v)
        
        max_depth = max(depth_nodes.keys()) if depth_nodes else 0
        
        # y 좌표 계산
        y_gap = (SCREEN_HEIGHT - 300) / (max_depth + 1) if max_depth > 0 else SCREEN_HEIGHT - 300
        
        # 각 깊이별로 노드 배치
        for depth, nodes in depth_nodes.items():
            num_nodes = len(nodes)
            x_gap = (SCREEN_WIDTH - 200) / (num_nodes + 1)
            
            for idx, node in enumerate(nodes):
                x = 100 + x_gap * (idx + 1)
                y = 100 + y_gap * depth
                self.positions[node] = (int(x), int(y))
    
    def draw_edge(self, u, v, color, width, highlight=False):
        # 간선 그리기
        if u not in self.positions or v not in self.positions:
            return
        
        pos_u = self.positions[u]
        pos_v = self.positions[v]
        
        if highlight:
            # 하이라이트 효과
            highlight_color = HIGHLIGHT_COLOR_QUERY if self.highlight_type == 'QUERY' else HIGHLIGHT_COLOR_UPDATE
            pygame.draw.line(self.screen, highlight_color, pos_u, pos_v, width + 6)
        
        pygame.draw.line(self.screen, color, pos_u, pos_v, width)
    
    def draw_node(self, v, color, highlight=False):
        # 노드 그리기
        if v not in self.positions:
            return
        
        pos = self.positions[v]
        
        # 하이라이트 효과
        if highlight:
            highlight_color = HIGHLIGHT_COLOR_QUERY if self.highlight_type == 'QUERY' else HIGHLIGHT_COLOR_UPDATE
            pygame.draw.circle(self.screen, highlight_color, pos, NODE_RADIUS + 7)
        
        # 노드 외곽선
        pygame.draw.circle(self.screen, BLACK, pos, NODE_RADIUS + 2)
        # 노드 내부
        pygame.draw.circle(self.screen, color, pos, NODE_RADIUS)
        
        # 노드 값 표시
        value_text = self.value_font.render(str(self.tree.val[v]), True, WHITE)
        value_rect = value_text.get_rect(center=pos)
        self.screen.blit(value_text, value_rect)
        
        # 노드 번호는 작게 위에 표시
        label_text = self.font.render(f"#{v}", True, BLACK)
        label_pos = (pos[0] - 15, pos[1] - NODE_RADIUS - 20)
        self.screen.blit(label_text, label_pos)
    
    def draw_tree(self):
        # 전체 트리 그리기

        # 간선 먼저 그리기
        for v in range(1, self.tree.n + 1):
            parent = self.tree.par[v]
            if parent != 0:
                edge_color = LIGHT_COLOR
                edge_width = EDGE_WIDTH_LIGHT
                
                # HLD의 경우 Heavy edge 강조
                if self.is_hld and self.tree.is_heavy(parent, v):
                    chain_id = self.tree.id[v]
                    color_idx = (chain_id - 1) % len(CHAIN_COLORS)
                    edge_color = CHAIN_COLORS[color_idx]
                    edge_width = EDGE_WIDTH_HEAVY
                
                # 간선 하이라이트: 양쪽 노드가 모두 하이라이트되어 있을 때만
                highlight = v in self.highlight_nodes and parent in self.highlight_nodes
                self.draw_edge(parent, v, edge_color, edge_width, highlight)
        
        # 노드 그리기
        for v in range(1, self.tree.n + 1):
            node_color = BLUE
            
            # HLD의 경우 체인별 색상
            if self.is_hld:
                chain_id = self.tree.id[v]
                color_idx = (chain_id - 1) % len(CHAIN_COLORS)
                node_color = CHAIN_COLORS[color_idx]
            
            highlight = v in self.highlight_nodes
            self.draw_node(v, node_color, highlight)
    
    def draw_info(self, selected_node):
        # 노드 정보 표시
        if selected_node == 0 or selected_node not in self.positions:
            return
        
        # 정보 패널 배경
        info_panel = pygame.Rect(10, SCREEN_HEIGHT - 250, 400, 240)
        pygame.draw.rect(self.screen, (240, 240, 240), info_panel)
        pygame.draw.rect(self.screen, BLACK, info_panel, 2)
        
        # 기본 정보
        info = [
            f"선택된 노드: #{selected_node}",
            f"노드 값: {self.tree.val[selected_node]}",
            f"깊이 (depth): {self.tree.dep[selected_node]}",
            f"부모 (parent): #{self.tree.par[selected_node]}",
            f"ETT in: {self.tree.in_time[selected_node]}",
            f"ETT out: {self.tree.out_time[selected_node]}"
        ]
        
        # HLD 추가 정보
        if self.is_hld:
            info.append(f"체인 ID: {self.tree.id[selected_node]}")
            info.append(f"체인 Head: #{self.tree.top[selected_node]}")
            info.append(f"서브트리 크기: {self.tree.size[selected_node]}")
        
        y_offset = SCREEN_HEIGHT - 240
        for i, text_str in enumerate(info):
            text = self.font.render(text_str, True, BLACK)
            self.screen.blit(text, (20, y_offset + i * 25))
    
    def draw_query_panel(self, query_mode, node1, node2, query_val):
        # 쿼리 패넗
        panel = pygame.Rect(SCREEN_WIDTH - 510, SCREEN_HEIGHT - 250, 500, 240)
        pygame.draw.rect(self.screen, (240, 240, 240), panel)
        pygame.draw.rect(self.screen, BLACK, panel, 2)
        
        y_offset = SCREEN_HEIGHT - 240
        
        # 제목
        title_text = "======= 경로 쿼리 =======" if self.is_hld else "====== 서브트리 쿼리 ======"
        title = self.font.render(title_text, True, BLACK)
        self.screen.blit(title, (SCREEN_WIDTH - 500, y_offset))
        y_offset += 30
        

        if query_mode == 'SELECT_FIRST':
            if self.is_hld:
                text = self.font.render("첫 번째 노드를 선택하세요", True, RED)
            else:
                text = self.font.render("서브트리의 루트 노드를 선택하세요", True, RED)
            self.screen.blit(text, (SCREEN_WIDTH - 500, y_offset))
        
        elif query_mode == 'SELECT_SECOND':
            text1 = self.font.render(f"첫 번째 노드: #{node1}", True, BLACK)
            self.screen.blit(text1, (SCREEN_WIDTH - 500, y_offset))
            y_offset += 25
            text2 = self.font.render("두 번째 노드를 선택하세요", True, RED)
            self.screen.blit(text2, (SCREEN_WIDTH - 500, y_offset))
        
        elif query_mode == 'INPUT_VALUE':
            if self.is_hld:
                text1 = self.font.render(f"노드: #{node1} -> #{node2}", True, BLACK)
            else:
                text1 = self.font.render(f"서브트리 루트: #{node1}", True, BLACK)
            self.screen.blit(text1, (SCREEN_WIDTH - 500, y_offset))
            y_offset += 25
            text2 = self.font.render(f"입력값: {query_val if query_val else ''}", True, BLUE)
            self.screen.blit(text2, (SCREEN_WIDTH - 500, y_offset))
            y_offset += 25
            text3 = self.font.render("값을 입력하고 Enter를 누르세요", True, RED)
            self.screen.blit(text3, (SCREEN_WIDTH - 500, y_offset))
        
        elif query_mode == 'EXECUTING':
            text1 = self.font.render(f"쿼리 실행 중...", True, GREEN)
            self.screen.blit(text1, (SCREEN_WIDTH - 500, y_offset))
            y_offset += 25
            if self.query_result is not None:
                text2 = self.font.render(f"결과: {self.query_result}", True, BLUE)
                self.screen.blit(text2, (SCREEN_WIDTH - 500, y_offset))
                y_offset += 25
            text3 = self.font.render("Enter: 완료", True, BLACK)
            self.screen.blit(text3, (SCREEN_WIDTH - 500, y_offset))
        
        # 단축키 안내
        y_offset = SCREEN_HEIGHT - 80
        text = self.font.render("ESC: 모드 취소", True, (100, 100, 100))
        self.screen.blit(text, (SCREEN_WIDTH - 500, y_offset))
    

    def draw_title(self):
        #제목
        title = "Heavy-Light Decomposition" if self.is_hld else "Euler Tour Technique"
        text = self.title_font.render(title, True, BLACK)
        text_rect = text.get_rect(center=(SCREEN_WIDTH // 2, 30))
        self.screen.blit(text, text_rect)
        
        # 조작 안내
        guide = "클릭: 노드 선택 | U: 업데이트 | Q: 쿼리 | ESC: 종료"
        guide_text = self.font.render(guide, True, (100, 100, 100))
        guide_rect = guide_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT - 20))
        self.screen.blit(guide_text, guide_rect)
    
    def find_clicked_node(self, mouse_pos):
        # 클린 노드 찾기
        min_dist = float('inf')
        closest_node = 0
        
        for node, pos in self.positions.items():
            dist_sq = (mouse_pos[0] - pos[0])**2 + (mouse_pos[1] - pos[1])**2
            if dist_sq < min_dist:
                min_dist = dist_sq
                closest_node = node
        
        # 노드 반경 내에 클릭했는지 확인
        if min_dist < (NODE_RADIUS * 2)**2:
            return closest_node
        return 0
    
    def set_highlight(self, nodes, op_type='QUERY'):
        # 하이라이트
        self.highlight_nodes = set(nodes) if nodes else set()
        self.highlight_type = op_type
    
    def run(self):
        running = True
        selected_node = 1
        clock = pygame.time.Clock()
        
        # 쿼리 모드 상태
        query_mode = None  # None, 'UPDATE', 'QUERY'
        query_state = 'IDLE'  # 'IDLE', 'SELECT_FIRST', 'SELECT_SECOND', 'INPUT_VALUE', 'EXECUTING'
        query_node1 = 0
        query_node2 = 0
        query_value = ""
        query_generator = None
        
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        if query_state != 'IDLE':
                            # 쿼리 모드 취소
                            query_mode = None
                            query_state = 'IDLE'
                            query_node1 = 0
                            query_node2 = 0
                            query_value = ""
                            self.set_highlight([])
                        else:
                            running = False
                    
                    # 쿼리 모드 시작
                    elif event.key == pygame.K_u and query_state == 'IDLE':
                        query_mode = 'UPDATE'
                        query_state = 'SELECT_FIRST'
                        query_value = ""
                    
                    elif event.key == pygame.K_q and query_state == 'IDLE':
                        query_mode = 'QUERY'
                        query_state = 'SELECT_FIRST'
                    
                    # 값 입력 모드
                    elif query_state == 'INPUT_VALUE':
                        if event.key == pygame.K_RETURN:
                            if query_value:
                                val = int(query_value)
                                if self.is_hld:
                                    if query_mode == 'UPDATE':
                                        query_generator = self.tree.update(query_node1, query_node2, val)
                                    else:
                                        query_generator = self.tree.query(query_node1, query_node2)
                                else:
                                    # ETT는 단일 노드만 사용
                                    if query_mode == 'UPDATE':
                                        query_generator = self.tree.update(query_node1, val)
                                    else:
                                        query_generator = self.tree.query(query_node1)
                                query_state = 'EXECUTING'
                                self.query_result = None
                        elif event.key == pygame.K_BACKSPACE:
                            query_value = query_value[:-1]
                        elif event.unicode.isdigit() or (event.unicode == '-' and not query_value):
                            query_value += event.unicode
                    # 실행중
                    elif query_state == 'EXECUTING' and event.key == pygame.K_RETURN:
                        if query_generator:
                            try:
                                nodes, op_type, result = next(query_generator)
                                self.set_highlight(nodes, op_type)
                                self.query_result = result
                            except StopIteration:
                                # 쿼리 완료
                                query_mode = None
                                query_state = 'IDLE'
                                query_generator = None
                                self.set_highlight([])
                                self.query_result = None
                
                # 노드 클릭
                if event.type == pygame.MOUSEBUTTONDOWN:
                    mouse_pos = pygame.mouse.get_pos()
                    clicked_node = self.find_clicked_node(mouse_pos)
                    
                    if clicked_node > 0:
                        if query_state == 'SELECT_FIRST':
                            query_node1 = clicked_node
                            # ETT는 단일 노드만 필요
                            if self.is_hld and (query_mode == 'QUERY' or query_mode == 'UPDATE'):
                                query_state = 'SELECT_SECOND'
                            else:
                                # ETT는 바로 값 입력 또는 쿼리 실행
                                if query_mode == 'UPDATE':
                                    query_state = 'INPUT_VALUE'
                                else:
                                    # 쿼리는 바로 실행
                                    query_generator = self.tree.query(query_node1)
                                    query_state = 'EXECUTING'
                            selected_node = clicked_node
                        elif query_state == 'SELECT_SECOND':
                            query_node2 = clicked_node
                            if query_mode == 'UPDATE':
                                query_state = 'INPUT_VALUE'
                            else:
                                # 쿼리는 바로 실행
                                if self.is_hld:
                                    query_generator = self.tree.query(query_node1, query_node2)
                                query_state = 'EXECUTING'
                            selected_node = clicked_node
                        else:
                            selected_node = clicked_node
            
            # 화면 그리기
            self.screen.fill(WHITE)
            self.draw_title()
            self.draw_tree()
            self.draw_info(selected_node)
            
            if query_state != 'IDLE':
                self.draw_query_panel(query_state, query_node1, query_node2, query_value)
            
            pygame.display.flip()
            clock.tick(60)
        
        pygame.quit()

render.py

 


4. 배운점, 느낀점 및 개선점

HLD와 ETT를 처음으로 배울때도 너무 재밌는 아이디어라 생각했는데,

이걸 직접 시각화해보고 프로젝트로 옮겨보니 원리에 대해 더 잘 이해할 수 있는 계기가 되었던 것 같다.

 

쿼리를 진행하는 트리의 구조를 하드코딩하는게 아니라 사용자가 트리를 입력하거나 실행 중 구조를 변경하게 하면 좋을 것 같다.

DFS ordering이나 ETT의 과정도 더 구체적인 단계로 보여주면 좋을 것 같다.


 

HLD와 ETT 구현에 대해 설명이 더 필요하다면

2025.10.24 - [PS] - [백준] 13510 트리와 쿼리 1 - C++

이 글을 참고해주세요..

 

파이게임 창이 화면캡쳐가 잘 안되어서 시연 영상은 없습니다.. 직접 해보세요!


'Unifox' 카테고리의 다른 글

[Unifox] - 웹해킹 2차시  (0) 2025.11.07
[Unifox] - 웹해킹 1차시  (0) 2025.11.02
[Unifox] - SQL  (1) 2025.10.03
[Unifox] - Flask 2차시  (0) 2025.10.03
[Unifox] - Flask 1차시  (0) 2025.09.22