ミラティブ プログラミングコンテスト2025（AtCoder Beginner Contest 414）E,F,G問題メモ

def solve(n):
    MOD = 998244353
    ans = (n + 1) * n
    ans -= (n + 1) * n // 2
    ans %= MOD

    tmp = 0
    for i in range(1, n + 1):
        if i * i > n:
            break
        tmp += 2 * (n // i)
    i -= 1
    tmp -= i * i
    ans -= tmp
    ans %= MOD

    return ans


n = int(input())
ans = solve(n)
print(ans)

各頂点 $(1,0)～(N,K-1)$ を $0～NK-1$ の整数値に対応づける方法で実装してしまったが、 直前の頂点 $p$ を判定するときに毎回割り算が必要になり、却って遅くなる実装だったかもしれない。
PyPyでもTLEが取れず、生成AIでRustに変換すると通った。

import sys
from collections import deque


def solve(n, k, edges):
    links = [[] for _ in range(n * k)]
    for u, v in edges:
        u -= 1
        v -= 1
        for j in range(k):
            links[j * n + u].append(((j + 1) % k) * n + v)
            links[j * n + v].append(((j + 1) % k) * n + u)

    INF = 1 << 60
    dist1 = [INF] * (n * k)
    dist2 = [INF] * (n * k)
    prev = [-1] * (n * k)
    dist1[0] = 0
    dist2[0] = 0
    q = deque([(0, -1, True)])
    while q:
        ku, p, is_first = q.popleft()
        j, u = divmod(ku, n)

        if is_first:
            dv = dist1[ku] + 1
            for kv in links[ku]:
                if ku >= n and p == kv % n:
                    continue
                if dist1[kv] == INF:
                    dist1[kv] = dv
                    prev[kv] = u
                    q.append((kv, u, True))
                elif dist2[kv] == INF and prev[kv] != u:
                    dist2[kv] = dv
                    q.append((kv, u, False))
        else:
            dv = dist2[ku] + 1
            kv = ((j + 1) % k) * n + prev[ku]
            if dist1[kv] == INF:
                dist1[kv] = dv
                prev[kv] = u
                q.append((kv, u, True))
            elif dist2[kv] == INF and prev[kv] != u:
                dist2[kv] = dv
                q.append((kv, u, False))

    ans = [-1] * n
    for v in range(1, n):
        if dist1[v] < INF:
            ans[v] = dist1[v] // k
    return ans[1:]


input = sys.stdin.buffer.readline
t = int(input())
for _ in range(t):
    n, k = map(int, input().split())
    edges = [tuple(map(int, input().split())) for _ in range(n - 1)]
    ans = solve(n, k, edges)
    print(*ans)

from heapq import heappop, heappush


def get_seg_nodes(a, b, offset):
    """ 全体 [0, offset) のうち、区間 [a, b) を被覆するセグメント木のノード番号のリストを得る """
    l = a + offset
    r = b + offset
    result_l = []
    result_r = []
    while l < r:
        if r & 1:
            result_r.append(r - 1)
        if l & 1:
            result_l.append(l)
            l += 1
        l >>= 1
        r >>= 1
    result_l.extend(reversed(result_r))
    return result_l


def get_index_range(k, depth):
    """ 深さ depth のセグメント木の、ノード番号 k が表す区間 [l, r) を得る """
    d = k.bit_length() - 1
    pos = k - (1 << d)
    width = 1 << (depth - d)
    return pos * width, (pos + 1) * width


def dijkstra(n, links, s=0):
    INF = 1 << 60
    dp = [INF] * n
    dp[s] = 0
    hq = [(0, s)]
    while hq:
        cost, u = heappop(hq)
        if dp[u] < cost:
            continue
        for v, c in links[u]:
            nc = cost + c
            if dp[v] <= nc:
                continue
            dp[v] = nc
            heappush(hq, (nc, v))
    return dp


def solve(n, m, xxx, expresses):
    INF = 1 << 60
    depth = (n - 1).bit_length()
    offset = 1 << depth
    ddd = [0] + [x1 - x0 for x0, x1 in zip(xxx, xxx[1:])]
    yyy = [0] + xxx  # 累積和計算用
    pre1 = [0] * (offset << 1)
    pre2 = [0] * (offset << 1)
    for i in range(n):
        pre1[i + offset] = ddd[i]
        pre2[i + offset - 1] = ddd[i]
    for i in range(offset - 1, 0, -1):
        pre1[i] = pre1[i * 2] + pre1[i * 2 + 1]
        pre2[i] = pre2[i * 2] + pre2[i * 2 + 1]

    os1 = offset
    os2 = offset * 2
    os3 = offset * 3
    os4 = offset * 4
    os5 = offset * 5

    gn = offset * 5 + m
    links = [[] for _ in range(gn)]
    for i in range(offset - 1, 0, -1):
        li = i * 2
        ri = i * 2 + 1

        # → 移動の乗り口
        if li < offset:
            lc = li + os1
            rc = ri + os1
        else:
            lc = li - offset
            rc = ri - offset
        links[lc].append((i + os1, pre1[ri]))
        links[rc].append((i + os1, 0))

        # ← 移動の乗り口
        if li < offset:
            lc = li + os2
            rc = ri + os2
        links[lc].append((i + os2, 0))
        links[rc].append((i + os2, pre2[li]))

        # → 移動の降り口
        if li < offset:
            lc = li + os3
            rc = ri + os3
        links[i + os3].append((lc, 0))
        links[i + os3].append((rc, pre2[li]))

        # ← 移動の降り口
        if li < offset:
            lc = li + os4
            rc = ri + os4
        links[i + os4].append((lc, pre1[ri]))
        links[i + os4].append((rc, 0))

    for i in range(m):
        l1, r1, l2, r2, c = expresses[i]
        l1 -= 1
        l2 -= 1
        e = i + os5

        seg1 = get_seg_nodes(l1, r1, offset)
        seg2 = get_seg_nodes(l2, r2, offset)
        # print(f'{l1=} {r1=} {l2=} {r2=} {c=}')
        # print(f'{seg1=} {seg2=}')
        if r1 <= l2:
            for k in seg1:
                if k >= offset:
                    k -= offset
                    links[k].append((e, yyy[r1] - xxx[k]))
                else:
                    s, t = get_index_range(k, depth)
                    links[k + os1].append((e, yyy[r1] - yyy[t]))
            for k in seg2:
                if k >= offset:
                    k -= offset
                    links[e].append((k, c + xxx[k] - yyy[r1]))
                else:
                    s, t = get_index_range(k, depth)
                    links[e].append((k + os3, c + xxx[s] - yyy[r1]))
        else:
            for k in seg1:
                if k >= offset:
                    k -= offset
                    links[k].append((e, xxx[k] - xxx[l1]))
                else:
                    s, t = get_index_range(k, depth)
                    links[k + os2].append((e, xxx[s] - xxx[l1]))
            for k in seg2:
                if k >= offset:
                    k -= offset
                    links[e].append((k, c + xxx[l1] - xxx[k]))
                else:
                    s, t = get_index_range(k, depth)
                    links[e].append((k + os4, c + xxx[l1] - yyy[t]))

    result = dijkstra(gn, links, 0)
    ans = []
    for i in range(1, n):
        if result[i] == INF:
            ans.append(-1)
        else:
            ans.append(result[i])
    return ans


n, m = map(int, input().split())
xxx = list(map(int, input().split()))
expresses = [tuple(map(int, input().split())) for _ in range(m)]
ans = solve(n, m, xxx, expresses)
print(*ans)