ALDS 1_7_A のC++のコードを解読しRustで再実装する


ALDS 1_7_A

https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ALDS1_7_A

根付き木の表現

  • 左子右兄弟表現
    • leftは自分の子の中で最も左のを表す
    • rightは自分の右隣の 兄弟を表す。
      従って、leftの深さは自分より1多く、rightの深さは同じ。

C++のコード

プログラミングコンテスト攻略のためのアルゴリズムとデータ構造に適宜コメントを加えたものです。

#include<iostream>
using namespace std;
#define MAX 100005
#define NIL -1

struct Node {int p, l, r ;};
// pは自分の親のノード
// lは自分から見て最も左の子のノード
// rは自分の右隣のノード

Node T[MAX];
int n, D[MAX];

void print(int u) {
    int i, c;
    cout << "node " << u << ": ";
    cout << "parent = " << T[u].p << ", ";
    cout << "depth = " << D[u] << ", ";

    if (T[u].p == NIL) cout << "root, " ;
    else if (T[u].l == NIL) cout << "leaf, "; 
    else cout << "internal node, ";

    cout << "[";

    for (i = 0, c= T[u].l; c != NIL; i++, c=T[c].r){
        if (i) cout << ", ";
        cout << c;
    }

    cout << "]" << endl ;
}

void rec (int u, int p){
    // u : ノード番号
    // p : 深さ
    // recはD[u]を変更する。
    D[u] = p; //u番目のノードの深さにpを設定

     // 右のノード => 同じ深さを設定
    if (T[u].r != NIL) rec(T[u].r, p);
    // T[u]が最も左のノードだった場合、次の段(p+1)に移動。
    if (T[u].l != NIL) rec(T[u].l, p+1);
}

int main(){
    int i,j,d,v,c,l,r;
    cin >> n;
    // T[]の初期化
    for (i=0; i<n;i++) T[i].p = T[i].l = T[i].r = NIL;

    for (i=0; i<n; i++){
        // vはid、dは子の数に対応
        cin >> v >> d;
        // id: vのノードは、d個の子ノードをもつ。
        for (j=0; j<d; j++){
            cin >> c;
            if (j==0) T[v].l = c; // 最初の子供をv番目の左側の子供として設定。
            else T[l].r = c; // l番目のノードの右隣ノードとしてノードcを設定
            l = c; // lにノードcを設定。次のノードの右隣ノードを決定するために使う。
            T[c].p = v; // c番目の親にid: vのノードを設定。
        }
    }
    for (i=0; i < n; i++){
        if (T[i].p==NIL) r=i; // 親がいないノードを探してrとしている。
    }
    // 親番号から探索開始
    rec(r,0);
    for (i=0; i<n; i++) print(i);
    return 0;
}

Rustでの再実装

i32やらusizeやらが混じっていてasが汚いのと、for文のparameter定義が汚いので、気が向いたら書き直す。

use std::io;

fn read<T: std::str::FromStr>() -> Vec<T> {
    let mut buf = String::new();
    io::stdin().read_line(&mut buf).unwrap();
    buf.trim().split(' ').flat_map(str::parse).collect()
}

#[derive(Debug, Clone)]
struct Node {
    parent: i32,
    left: i32,
    right: i32
}

impl Node {
    fn set_left(&mut self, left: i32) {
        self.left = left
    }
    fn set_right(&mut self, right: i32) {
        self.right = right
    }
}

fn printer(id :usize, tree: &Vec<Node>, depth: &Vec<i32>){
    let node_type = if tree[id].parent == -1 {
        "root"
    } else if tree[id].left == -1 {
        "leaf"
    } else {
        "internal node"
    };

    print!("node {node}: parent = {parent}, depth = {depth}, {node_type}, [",
    node = id, 
    parent = tree[id].parent, 
    depth = depth[id], 
    node_type = node_type);
    let mut i = 0;
    let mut c = tree[id].left;
    while c != -1 { // leftがいるとき
        if i > 0 {print!(", ")};
        print!("{}", c);
        i += 1;
        c=tree[c as usize].right;
    }
    println!("]")


}

fn calc_depth(tree: &Vec<Node>, depth: &mut Vec<i32>, id: usize, d: usize){
    depth[id] = d as i32;

    if tree[id].right != -1 {calc_depth(&tree, depth, tree[id].right as usize, d)}
    if tree[id].left != -1 {calc_depth(&tree, depth, tree[id].left as usize, d+1)}
}

fn main() {
    let node_num = read::<usize>()[0];
    // T[]に対応
    let mut tree = vec![Node{parent: -1,left: -1, right: -1}; node_num];
    // D[]に対応
    let mut depth = vec![0; node_num];

    
    for _i in 0..node_num {
        let mut line = read::<usize>();
        let id = line.remove(0);
        let _child_num = line.remove(0) as usize;
        let childs = line;
        let mut counter = 0;
        let mut left = 0;
        for c in childs {
            if counter == 0 {
                tree[id].set_left(c as i32);
            } else {
                tree[left].set_right(c as i32);
            }
            left = c as usize;
            tree[c].parent = id as i32;
            counter += 1
        }
        
    }
    let mut root = 0;
    for i in 0..node_num {
        if tree[i].parent == -1 {
            root = i;
        }
    }
    calc_depth(&tree , &mut depth, root, 0);
    for i in 0..node_num {printer(i, &tree, &depth)}
}