#1159590 (Rust) No.3506 All Distance is Square Number

• 問題: No.3506 All Distance is Square Number
• ユーザー: ガルム
• 提出日時: 2026-04-18 14:02:35
• 言語: Rust (1.94.0 + proconio + num + itertools)
• 結果: AC
• 実行時間: 315 ms / 2,000 ms
• コード長: 7,540 bytes
• コンパイル時間: 4,474 ms
• コンパイル使用メモリ: 200,416 KB
• 実行使用メモリ: 6,400 KB
• 最終ジャッジ日時: 2026-04-18 14:02:45
• 合計ジャッジ時間: 6,898 ms
• ジャッジサーバーID （参考情報）: judge3_0 / judge1_1

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

ソースコード

use std::io::{self, Read};

#[derive(Clone)]
struct Edge {
    u: usize,
    v: usize,
    w: usize,
}

fn has(bs: &[u64], x: usize) -> bool {
    ((bs[x >> 6] >> (x & 63)) & 1) == 1
}

fn set_bit(bs: &mut [u64], x: usize) {
    bs[x >> 6] |= 1u64 << (x & 63);
}

fn shift_or(dst: &mut [u64], src: &[u64], sh: usize, max_sum: usize) {
    let q = sh >> 6;
    let r = sh & 63;
    let words = dst.len();

    for i in 0..words {
        let x = src[i];
        if x == 0 {
            continue;
        }
        if i + q < words {
            dst[i + q] |= x << r;
        }
        if r != 0 && i + q + 1 < words {
            dst[i + q + 1] |= x >> (64 - r);
        }
    }

    let rem = (max_sum + 1) & 63;
    if rem != 0 {
        let mask = (1u64 << rem) - 1;
        let last = words - 1;
        dst[last] &= mask;
    }
}

fn arc_positions(n: usize, from: usize, to: usize) -> Vec<usize> {
    let mut res = Vec::new();
    let mut x = from;
    while x != to {
        res.push(x);
        x += 1;
        if x == n + 1 {
            x = 1;
        }
    }
    res
}

fn arc_path(
    n: usize,
    by_pos: &[Vec<(usize, usize)>],
    sqs: &[usize],
    from: usize,
    to: usize,
    max_sum: usize,
) -> Option<Vec<usize>> {
    let pos = arc_positions(n, from, to);
    let words = (max_sum + 64) >> 6;
    let m = pos.len();
    let mut suf = vec![vec![0u64; words]; m + 1];
    set_bit(&mut suf[m], 0);

    for i in (0..m).rev() {
        let p = pos[i];
        let mut cur = vec![0u64; words];
        for &(_, w) in &by_pos[p] {
            shift_or(&mut cur, &suf[i + 1], w, max_sum);
        }
        suf[i] = cur;
    }

    let mut target = None;
    for &x in sqs {
        if x <= max_sum && has(&suf[0], x) {
            target = Some(x);
            break;
        }
    }
    let mut rem = target?;
    let mut path = Vec::new();

    for i in 0..m {
        let p = pos[i];
        for &(id, w) in &by_pos[p] {
            if rem >= w && has(&suf[i + 1], rem - w) {
                path.push(id);
                rem -= w;
                break;
            }
        }
    }

    Some(path)
}

fn dfs_small(
    v: usize,
    t: usize,
    sum: usize,
    adj: &[Vec<(usize, usize, usize)>],
    sq: &[bool],
    used: &mut [bool],
    path: &mut Vec<usize>,
) -> bool {
    if v == t {
        return sq[sum];
    }

    for &(to, id, w) in &adj[v] {
        if !used[to] {
            used[to] = true;
            path.push(id);
            if dfs_small(to, t, sum + w, adj, sq, used, path) {
                return true;
            }
            path.pop();
            used[to] = false;
        }
    }

    false
}

fn small_graph(n: usize) -> Vec<Edge> {
    let uv = [
        (1, 2),
        (1, 3),
        (2, 3),
        (1, 4),
        (2, 4),
        (1, 5),
        (2, 5),
        (1, 6),
        (2, 6),
        (1, 7),
        (6, 7),
        (1, 8),
        (7, 8),
        (3, 9),
        (7, 9),
        (1, 10),
        (7, 10),
        (1, 11),
        (3, 11),
        (1, 12),
        (3, 12),
        (1, 13),
        (3, 13),
        (1, 14),
        (2, 14),
        (1, 15),
        (2, 15),
        (1, 16),
        (3, 16),
    ];
    let ww = [
        4, 36, 196, 25, 169, 144, 9, 16, 7, 24, 17, 2, 1, 5, 3, 6, 8, 10, 11,
        12, 13, 14, 15, 18, 19, 21, 20, 22, 23,
    ];

    let m = 2 * n - 3;
    let mut edges = Vec::new();
    for i in 0..m {
        edges.push(Edge {
            u: uv[i].0,
            v: uv[i].1,
            w: ww[i],
        });
    }
    edges
}

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

    let mut edges = Vec::new();
    let mut by_pos = vec![Vec::new(); n + 1];
    let mut idx = 0usize;

    for p in 1..=n {
        let u = p;
        let v = if p == n { 1 } else { p + 1 };
        let id = edges.len() + 1;
        let w = ws[idx];
        idx += 1;
        edges.push(Edge { u, v, w });
        by_pos[p].push((id, w));

        if p <= n - 3 {
            let id = edges.len() + 1;
            let w = ws[idx];
            idx += 1;
            edges.push(Edge { u, v, w });
            by_pos[p].push((id, w));
        }
    }

    (edges, by_pos)
}

fn main() {
    let mut input = String::new();
    io::stdin().read_to_string(&mut input).unwrap();
    let n: usize = input.trim().parse().unwrap();

    let mut out = String::new();

    if n <= 16 {
        let edges = small_graph(n);
        let max_sum: usize = edges.iter().map(|e| e.w).sum();
        let mut sq = vec![false; max_sum + 1];
        let mut x = 1usize;
        while x * x <= max_sum {
            sq[x * x] = true;
            x += 1;
        }

        let mut adj = vec![Vec::new(); n + 1];
        for (i, e) in edges.iter().enumerate() {
            let id = i + 1;
            adj[e.u].push((e.v, id, e.w));
            adj[e.v].push((e.u, id, e.w));
        }

        out.push_str(&format!("{}\n", edges.len()));
        for e in &edges {
            out.push_str(&format!("{} {} {}\n", e.u, e.v, e.w));
        }

        for i in 1..=n {
            for j in i + 1..=n {
                let mut used = vec![false; n + 1];
                let mut path = Vec::new();
                used[i] = true;
                dfs_small(i, j, 0, &adj, &sq, &mut used, &mut path);
                out.push_str(&path.len().to_string());
                for id in path {
                    out.push_str(&format!(" {}", id));
                }
                out.push('\n');
            }
        }
    } else {
        let (edges, by_pos) = cycle_graph(n);
        let max_sum: usize = edges.iter().map(|e| e.w).sum();
        let mut sqs = Vec::new();
        let mut x = 1usize;
        while x * x <= max_sum {
            sqs.push(x * x);
            x += 1;
        }

        out.push_str(&format!("{}\n", edges.len()));
        for e in &edges {
            out.push_str(&format!("{} {} {}\n", e.u, e.v, e.w));
        }

        for i in 1..=n {
            for j in i + 1..=n {
                let path = if let Some(p) = arc_path(n, &by_pos, &sqs, i, j, max_sum) {
                    p
                } else {
                    let mut p = arc_path(n, &by_pos, &sqs, j, i, max_sum).unwrap();
                    p.reverse();
                    p
                };
                out.push_str(&path.len().to_string());
                for id in path {
                    out.push_str(&format!(" {}", id));
                }
                out.push('\n');
            }
        }
    }

    print!("{}", out);
}