結果
問題 |
No.650 行列木クエリ
|
ユーザー |
![]() |
提出日時 | 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 |
ソースコード
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()