結果

問題 No.650 行列木クエリ
ユーザー qwewe
提出日時 2025-05-14 13:20:31
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 935 ms / 2,000 ms
コード長 6,205 bytes
コンパイル時間 389 ms
コンパイル使用メモリ 82,516 KB
実行使用メモリ 274,304 KB
最終ジャッジ日時 2025-05-14 13:22:30
合計ジャッジ時間 5,412 ms
ジャッジサーバーID
(参考情報)
judge2 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 1
other AC * 10
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys

def solve():
    input = sys.stdin.readline
    sys.setrecursionlimit(2 * 10**5) # Max N is 10^5, DFS could go that deep

    N = int(input())
    MOD = 10**9 + 7

    adj = [[] for _ in range(N)]
    for i in range(N - 1):
        u, v = map(int, input().split())
        adj[u].append((v, i)) # Store neighbor and edge index
        adj[v].append((u, i))

    parent = [-1] * N
    depth = [-1] * N 
    subtree_size = [0] * N
    # edge_idx_to_child_node[edge_k] = v means edge k connects parent[v] and v
    edge_idx_to_child_node = [-1] * (N - 1) 

    # HLD arrays
    heavy_child = [-1] * N 
    head = [-1] * N          # Head of the heavy path for a node
    pos = [-1] * N           # Position in the segment tree's base array
    hld_timer = 0
    
    identity_matrix = ((1, 0), (0, 1))

    def mat_mul(A, B):
        C00 = (A[0][0] * B[0][0] + A[0][1] * B[1][0]) % MOD
        C01 = (A[0][0] * B[0][1] + A[0][1] * B[1][1]) % MOD
        C10 = (A[1][0] * B[0][0] + A[1][1] * B[1][0]) % MOD
        C11 = (A[1][0] * B[0][1] + A[1][1] * B[1][1]) % MOD
        return ((C00, C01), (C10, C11))

    def dfs_size(u, p_node, d): # u: current, p_node: parent, d: depth
        parent[u] = p_node
        depth[u] = d
        subtree_size[u] = 1
        max_sz = 0
        
        for v, edge_i in adj[u]:
            if v == p_node:
                continue
            edge_idx_to_child_node[edge_i] = v # u is parent of v for this edge
            dfs_size(v, u, d + 1)
            subtree_size[u] += subtree_size[v]
            if subtree_size[v] > max_sz:
                max_sz = subtree_size[v]
                heavy_child[u] = v
    
    dfs_size(0, -1, 0) # Root is 0

    def dfs_hld(u, current_h): # u: current, current_h: head of heavy path for u
        nonlocal hld_timer
        head[u] = current_h
        pos[u] = hld_timer
        hld_timer += 1
        
        if heavy_child[u] != -1:
            dfs_hld(heavy_child[u], current_h) # Continue current heavy path
        
        for v, edge_i in adj[u]: # For light children
            if v == parent[u] or v == heavy_child[u]:
                continue
            dfs_hld(v, v) # v starts a new heavy path, so it's the head
    
    dfs_hld(0, 0) # Start HLD from root

    # Segment Tree
    # Uses 1-based indexing for tree nodes, covers HLD positions 0 to N-1.
    actual_seg_tree_array = [identity_matrix] * (4 * N)

    def build_st(st_idx, hld_l, hld_r): 
        if hld_l == hld_r:
            actual_seg_tree_array[st_idx] = identity_matrix
            return
        mid = (hld_l + hld_r) // 2
        build_st(2 * st_idx, hld_l, mid)
        build_st(2 * st_idx + 1, mid + 1, hld_r)
        actual_seg_tree_array[st_idx] = mat_mul(actual_seg_tree_array[2*st_idx], actual_seg_tree_array[2*st_idx+1])

    if N > 0: # Problem constraint N >= 2
      build_st(1, 0, N - 1)

    def update_st(st_idx, hld_l, hld_r, target_hld_pos, new_matrix):
        if hld_l == hld_r:
            actual_seg_tree_array[st_idx] = new_matrix
            return
        mid = (hld_l + hld_r) // 2
        if target_hld_pos <= mid:
            update_st(2 * st_idx, hld_l, mid, target_hld_pos, new_matrix)
        else:
            update_st(2 * st_idx + 1, mid + 1, hld_r, target_hld_pos, new_matrix)
        actual_seg_tree_array[st_idx] = mat_mul(actual_seg_tree_array[2*st_idx], actual_seg_tree_array[2*st_idx+1])

    def query_st(st_idx, hld_l, hld_r, query_l, query_r):
        if query_l > query_r: # Empty range
            return identity_matrix
        if query_l <= hld_l and hld_r <= query_r: # Current segment fully in query range
            return actual_seg_tree_array[st_idx]
        if hld_r < query_l or hld_l > query_r: # Current segment fully outside query range
            return identity_matrix
        
        mid = (hld_l + hld_r) // 2
        # Recursively query children and combine results
        # Note: The ranges passed to children are still [query_l, query_r].
        # Each child will determine its overlap.
        res_left = query_st(2 * st_idx, hld_l, mid, query_l, query_r)
        res_right = query_st(2 * st_idx + 1, mid + 1, hld_r, query_l, query_r)
        return mat_mul(res_left, res_right)

    Q = int(input())
    results_buffer = [] # Buffer output lines for faster printing

    for _ in range(Q):
        query_parts = input().split()
        type = query_parts[0]
        if type == 'x':
            edge_i = int(query_parts[1])
            vals = [int(x) for x in query_parts[2:]]
            new_mat = ((vals[0], vals[1]), (vals[2], vals[3]))
            
            # Matrix for edge (parent[v], v) is stored at pos[v]
            node_v_for_edge = edge_idx_to_child_node[edge_i]
            hld_pos_to_update = pos[node_v_for_edge]
            update_st(1, 0, N - 1, hld_pos_to_update, new_mat)
        else: # type == 'g'
            u, v = int(query_parts[1]), int(query_parts[2]) # u is ancestor of v
            
            res_matrix = identity_matrix
            curr_v_on_path = v 
            
            while head[curr_v_on_path] != head[u]:
                # Segment from head[curr_v_on_path] down to curr_v_on_path
                # HLD positions: pos[head[curr_v_on_path]] to pos[curr_v_on_path]
                segment_prod = query_st(1, 0, N - 1, pos[head[curr_v_on_path]], pos[curr_v_on_path])
                res_matrix = mat_mul(segment_prod, res_matrix) # Prepend product
                curr_v_on_path = parent[head[curr_v_on_path]]
            
            # curr_v_on_path and u are on the same heavy path. u is ancestor.
            # Path segment from u down to curr_v_on_path.
            # Matrices correspond to HLD positions pos[u]+1 to pos[curr_v_on_path].
            # If u == curr_v_on_path, range is empty (pos[u]+1 > pos[curr_v_on_path]), query_st returns identity.
            segment_prod = query_st(1, 0, N - 1, pos[u] + 1, pos[curr_v_on_path])
            res_matrix = mat_mul(segment_prod, res_matrix) # Prepend final segment product
            
            results_buffer.append(f"{res_matrix[0][0]} {res_matrix[0][1]} {res_matrix[1][0]} {res_matrix[1][1]}")

    sys.stdout.write("\n".join(results_buffer) + "\n")

solve()
0