結果

問題 No.2258 The Jikka Tree
ユーザー gew1fw
提出日時 2025-06-12 17:14:34
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 3,239 bytes
コンパイル時間 277 ms
コンパイル使用メモリ 82,352 KB
実行使用メモリ 85,904 KB
最終ジャッジ日時 2025-06-12 17:14:45
合計ジャッジ時間 9,512 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1
other AC * 2 WA * 1 TLE * 1 -- * 71
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
import bisect

sys.setrecursionlimit(1 << 25)

def main():
    input = sys.stdin.read
    data = input().split()
    ptr = 0

    N = int(data[ptr])
    ptr += 1

    # Read edges
    edges = [[] for _ in range(N)]
    for _ in range(N-1):
        u = int(data[ptr])
        v = int(data[ptr+1])
        ptr += 2
        edges[u].append(v)
        edges[v].append(u)

    # Compute in_order and out_order via DFS
    in_order = [0] * N
    out_order = [0] * N
    time = 0
    visited = [False] * N

    def dfs(u, parent):
        nonlocal time
        in_order[u] = time
        time += 1
        for v in edges[u]:
            if v != parent and not visited[v]:
                visited[v] = True
                dfs(v, u)
        out_order[u] = time - 1

    visited[0] = True
    dfs(0, -1)

    # Read A array
    A = list(map(int, data[ptr:ptr+N]))
    ptr += N

    # Prefix sum of A
    prefix = [0] * (N + 1)
    for i in range(N):
        prefix[i+1] = prefix[i] + A[i]

    # Read Q
    Q = int(data[ptr])
    ptr += 1

    queries = []
    for _ in range(Q):
        a_prime = int(data[ptr])
        b_prime = int(data[ptr+1])
        k_prime = int(data[ptr+2])
        delta = int(data[ptr+3])
        ptr += 4
        queries.append((a_prime, b_prime, k_prime, delta))

    # Process each query
    X = []
    sum_X = 0
    for i in range(Q):
        a_prime, b_prime, k_prime, delta = queries[i]

        # Compute a, b, k
        if i == 0:
            a = a_prime
            b = b_prime
            k = k_prime
        else:
            a = (a_prime + sum_X) % N
            b = (b_prime + 2 * sum_X) % N
            k = (k_prime + (sum_X ** 2)) % 150001

        # Compute l and r
        l = min(a, b)
        r = max(a, b) + 1

        # Compute T
        T = (prefix[r] - prefix[l]) + k * (r - l)

        # Function to count nodes in [l, r) with in_order in [a, b]
        def count(l, r, a, b):
            count = 0
            for w in range(l, r):
                if in_order[w] >= a and in_order[w] <= b:
                    count += 1
            return count

        # Function to sum A_w for nodes in [l, r) with in_order in [a, b]
        def sum_A(l, r, a, b):
            s = 0
            for w in range(l, r):
                if in_order[w] >= a and in_order[w] <= b:
                    s += A[w]
            return s

        # Find the optimal node v
        current_v = 0
        while True:
            max_sum = 0
            best_child = -1
            for v in edges[current_v]:
                if in_order[v] < in_order[current_v]:
                    # This is the parent, skip
                    continue
                # Compute sum for subtree of v
                cnt = count(l, r, in_order[v], out_order[v])
                s = sum_A(l, r, in_order[v], out_order[v])
                sum_u = cnt * k + s
                if sum_u > max_sum:
                    max_sum = sum_u
                    best_child = v
            if max_sum > T / 2:
                current_v = best_child
            else:
                break

        X.append(current_v)
        sum_X += current_v

    # Output the results
    for x in X:
        print(x)

if __name__ == '__main__':
    main()
0