#1158832 (PyPy3) No.3506 All Distance is Square Number

• 問題: No.3506 All Distance is Square Number
• ユーザー: Shiny
• 提出日時: 2026-04-18 01:31:36
• 言語: PyPy3 (7.3.17)
• 結果: AC
• 実行時間: 418 ms / 2,000 ms
• コード長: 5,687 bytes
• コンパイル時間: 325 ms
• コンパイル使用メモリ: 85,504 KB
• 実行使用メモリ: 131,988 KB
• 最終ジャッジ日時: 2026-04-18 01:31:48
• 合計ジャッジ時間: 8,145 ms
• ジャッジサーバーID （参考情報）: judge1_0 / judge2_0

ファイルパターン	結果
other	AC * 29

ソースコード

import sys
from functools import lru_cache

input = sys.stdin.readline

step1 = [
    1, 3, 4, 2, 7, 6, 9, 10, 11, 12, 13, 14, 17, 18, 19, 21, 24, 26, 28,
    32, 30, 29, 31, 34, 39, 37, 43, 47, 44, 46, 48, 55, 56, 58, 62, 59, 69,
    74, 63, 68, 72, 71, 84, 79, 80, 85, 83, 82, 88, 95, 96, 99, 91, 78, 105,
    93, 101, 110, 103, 107, 115, 119, 114, 122, 124, 126, 127, 129, 134, 135,
    128, 143, 146, 140, 139, 152, 138, 147, 155, 148, 159, 151, 156, 166, 160,
    164, 168, 170, 175, 172, 177, 178, 180, 181, 186, 192, 190, 187, 200,
]
step2 = [
    8, 25, 5, 42, 40, 15, 16, 22, 49, 20, 36, 50, 64, 23, 35, 27, 121, 144,
    45, 41, 33, 81, 57, 51, 38, 196, 67, 54, 53, 52, 60, 100, 61, 66, 65, 73,
    75, 77, 70, 169, 76, 87, 86, 98, 89, 92, 94, 90, 113, 97, 102, 106, 104,
    109, 111, 108, 117, 112, 116, 118, 120, 125, 123, 132, 133, 130, 131, 137,
    142, 136, 145, 150, 141, 149, 153, 154, 158, 157, 162, 167, 163, 165, 174,
    161, 171, 173, 179, 176, 182, 188, 184, 183, 185, 189, 193, 195, 191, 197,
]

N = int(input())
max_sum = 200 * N
square_list = []
now = 1
while now * now <= max_sum:
    square_list.append(now * now)
    now += 1

# snapshot arrays for the recurrence.
# P[t][i]  : sums of simple paths from i to t-1 in G_{t+1}
# Q[t][i]  : sums of simple paths from i to t   in G_{t+1}
# P0/Q0    : same, but avoiding the other end of the last edge pair
ps = [[0] * N for _ in range(N)]
qs = [[0] * N for _ in range(N)]
p0s = [[0] * N for _ in range(N)]
q0s = [[0] * N for _ in range(N)]

cur_p = [0] * N
cur_q = [0] * N
cur_p0 = [0] * N
cur_q0 = [0] * N

cur_p[0] = 1 << 0
cur_q[0] = 1 << step1[0]
cur_p0[0] = 1 << 0
cur_q0[0] = 0

cur_p[1] = 1 << step1[0]
cur_q[1] = 1 << 0
cur_p0[1] = 0
cur_q0[1] = 1 << 0

for i in range(2):
    ps[1][i] = cur_p[i]
    qs[1][i] = cur_q[i]
    p0s[1][i] = cur_p0[i]
    q0s[1][i] = cur_q0[i]

for t in range(2, N):
    add_last = step1[t - 1]
    add_skip = step2[t - 2]

    old_p = cur_p
    old_q = cur_q
    old_p0 = cur_p0

    cur_p = [0] * N
    cur_q = [0] * N
    cur_p0 = [0] * N
    cur_q0 = [0] * N

    join_cost = add_last + add_skip
    for i in range(t):
        cur_p[i] = old_q[i] | (old_p0[i] << join_cost)
        cur_p0[i] = old_q[i]
        cur_q[i] = (old_q[i] << add_last) | (old_p[i] << add_skip)
        cur_q0[i] = old_p0[i] << add_skip

    cur_p[t] = (1 << add_last) | (old_q[t - 2] << add_skip)
    cur_p0[t] = 0
    cur_q[t] = 1
    cur_q0[t] = 1

    for i in range(t + 1):
        ps[t][i] = cur_p[i]
        qs[t][i] = cur_q[i]
        p0s[t][i] = cur_p0[i]
        q0s[t][i] = cur_q0[i]

edges = []
edge_id = {}
for i in range(N - 1):
    idx = len(edges) + 1
    edges.append((i, i + 1, step1[i]))
    edge_id[(i, i + 1)] = idx
for i in range(N - 2):
    idx = len(edges) + 1
    edges.append((i, i + 2, step2[i]))
    edge_id[(i, i + 2)] = idx


def has_sum(bitset, value):
    return value >= 0 and ((bitset >> value) & 1)


def first_square(bitset):
    for sq in square_list:
        if has_sum(bitset, sq):
            return sq
    return -1


@lru_cache(None)
def restore_q(t, start, total):
    if start == t:
        return () if total == 0 else None
    if t == 1:
        if start == 0 and total == step1[0]:
            return (edge_id[(0, 1)],)
        return None

    add_last = step1[t - 1]
    add_skip = step2[t - 2]

    if start <= t - 1 and total >= add_last and has_sum(qs[t - 1][start], total - add_last):
        prev = restore_q(t - 1, start, total - add_last)
        if prev is not None:
            return prev + (edge_id[(t - 1, t)],)

    if start <= t - 1 and total >= add_skip and has_sum(ps[t - 1][start], total - add_skip):
        prev = restore_p(t - 1, start, total - add_skip)
        if prev is not None:
            return prev + (edge_id[(t - 2, t)],)

    return None


@lru_cache(None)
def restore_p(t, start, total):
    target = t - 1
    if start == target:
        return () if total == 0 else None
    if t == 1:
        if start == 1 and total == step1[0]:
            return (edge_id[(0, 1)],)
        if start == 0 and total == 0:
            return ()
        return None

    add_last = step1[t - 1]
    add_skip = step2[t - 2]

    if start < t and has_sum(qs[t - 1][start], total):
        prev = restore_q(t - 1, start, total)
        if prev is not None:
            return prev

    if start < t and total >= add_last + add_skip and has_sum(p0s[t - 1][start], total - add_last - add_skip):
        prev = restore_p0(t - 1, start, total - add_last - add_skip)
        if prev is not None:
            return prev + (edge_id[(t - 2, t)], edge_id[(t - 1, t)])

    if start == t:
        if total == add_last:
            return (edge_id[(t - 1, t)],)
        if total >= add_skip and has_sum(qs[t - 1][t - 2], total - add_skip):
            prev = restore_q(t - 1, t - 2, total - add_skip)
            if prev is not None:
                return (edge_id[(t - 2, t)],) + prev

    return None


@lru_cache(None)
def restore_p0(t, start, total):
    target = t - 1
    if start == target:
        return () if total == 0 else None
    if t == 1:
        if start == 0 and total == 0:
            return ()
        return None

    if start < t and has_sum(qs[t - 1][start], total):
        prev = restore_q(t - 1, start, total)
        if prev is not None:
            return prev
    return None


out = []
out.append(str(len(edges)))
for u, v, w in edges:
    out.append(f"{u + 1} {v + 1} {w}")

for i in range(N):
    for j in range(i + 1, N):
        wanted = first_square(qs[j][i])
        route = restore_q(j, i, wanted)
        out.append(str(len(route)) + " " + " ".join(map(str, route)))

sys.stdout.write("\n".join(out))