結果
| 問題 |
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 |
ソースコード
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()