2分木をASCIIで描くプログラム

SICPで2分木を書く問題が出てきましたが、ASCIIで木を手書きすると手間がかかるので、木を出力してくれるプログラムをSchemeで書いてみました。Gaucheで動作します。

こんな感じの表示になります。出力はあまりコンパクトにはなりません。(左右の子の木は列で重複しません)

     M        
    / \       
   /   \      
  C     \     
 / \     \    
*   E     \   
           S  
          / \ 
         /   *
        R     
       / \    
      *   *    

木の入力形式はS式で

葉 := () | ペア以外
木 := 葉 | (ラベル 左の木 右の木)

です。

binary-tree->stringは木を与えると文字列表現を買えし、print-binary-treeは出力します。

(porint-binary-tree '(A () B))

  A  
 / \ 
*   B
(print-binary-tree '(1 () (2 () (3 () (4 () (5 () (6 () (7 8 9))))))))

  1                    
 / \                   
*   \                  
     2                 
    / \                
   *   \               
        3              
       / \             
      *   \            
           4           
          / \          
         *   \         
              5        
             / \       
            *   \      
                 6     
                / \    
               *   \   
                    7  
                   / \ 
                  8   9

ソースコード

なんとなくスクリプトになってます。標準入力からS式の木を読み込んで出力をします。

#!/usr/bin/env gosh
;;
;; Draw Binary Tree
;; 
;; Author : Taisuke HORI
;; Date   : 2007/04/01
;;

(define (make-tree entry left right)
    (list entry left right))
(define (leaf? tree) (not (pair? tree)))
(define (entry tree) (car tree))
(define (entry->string e) (x->string e))
(define (left-branch tree) (cadr tree))
(define (right-branch tree) (caddr tree))

(define (make-rect strings n c)
  (list strings n c))
(define (rect-width r) (cadr r))
(define (rect-strings r) (car r))
(define (rect-center r) (caddr r))

;; エントリーの文字列と左右でtotal-widthをはみ出した長さを返す。
;; 文字列の長さがtotal-widthを越える場合のみ、はみ出した長さが0でなくなる。
(define (make-entry-string total-width center s)
  (let ((w (string-length s)))
    (if (< total-width w)
        (let* ((remain (- w total-width))
               (l-remain (quotient remain 2))
               (r-remain (- remain l-remain)))
          (values s l-remain r-remain))
        (let* ((lw (- center (quotient w 2)))
               (rw (- total-width w lw)))
          (when (< lw 0)
            (set! lw 0)
            (set! rw (- total-width w)))
          (when (< rw 0)
            (set! rw 0)
            (set! lw (- total-width w)))
          (values (string-append (make-string lw #\space)
                                 s
                                 (make-string rw #\space))
                  0 0)))))

(define (for-all proc . params)
  (if (null? params)
      (proc)
      (for-each (lambda (e)
                  (print "e: " e)
                  (apply (cut for-all (cut proc e <...>) <...>)
                         (cdr params)))
                (car params))))

;; sの左にleft-num, 右にright-numだけ空白を付ける。
(define (expand-string s left-num right-num )
  (string-append (make-string left-num #\space)
                 s
                 (make-string right-num #\space)))

;; リストがnより短ければ長さがnになるようにeを末尾に加える
(define (expand-list l n e)
  (let ((len (length l)))
    (if (< len n)
        (append l (make-list (- n len) e))
        l)))
            
(define (make-left-edge n)
  (let loop ((k 0))
    (if (= k n)
        '()
        (cons (string-append (make-string (- n k 1) #\space)
                             "/"
                             (make-string k #\space))
              (loop (+ k 1))))))

(define (make-right-edge n)
  (let loop ((k 0))
    (if (= k n)
        '()
        (cons (string-append (make-string k #\space)
                                     "\\"
                                     (make-string (- n k 1) #\space))
              (loop (+ k 1))))))

(define (append-space s l r)
  (string-append (make-string l #\space)
                 s
                 (make-string r #\space)))

(define (binary-tree->rect tree)
  (cond ;; 空の葉
        ((null? tree) (make-rect (list "*") 1 0))
        ;; 葉
        ((leaf? tree)
         (let* ((s (entry->string tree))
                (n (string-length s)))
           (make-rect (list s) n (quotient (- n 1) 2))))
        (else
         (let* ((lrect (binary-tree->rect (left-branch tree)))
                (rrect (binary-tree->rect (right-branch tree)))
                ;; 子の幅が1の場合枝がはみ出る
                (lprotrusion (if (= 1 (rect-width lrect)) 1 0))
                (rprotrusion (if (= 1 (rect-width rrect)) 1 0))
                (ls (map (cut append-space <> 0 lprotrusion) (rect-strings lrect)))
                (lw (+ lprotrusion (rect-width lrect)))
                (rs (map (cut append-space <> rprotrusion 0) (rect-strings rrect)))
                (rw (+ rprotrusion (rect-width rrect)))
                (ehl (- lw (rect-center lrect) 1)) ; 左の枝の長さ
                (ehr (max 1 (rect-center rrect))) ; 右の枝の長さ
                (tw (+ 1 lw rw))
                (ledges (map (lambda (e)
                               (append-space e (- lw (string-length e)) 0))
                             (make-left-edge ehl)))
                (redges (map (lambda (e)
                               (append-space e 0 (- rw (string-length e))))
                             (make-right-edge ehr)))
                ;; 左右の高さの最大
                (h (max (+ ehl (length ls)) (+ ehr (length rs))))
                (center (+ 1 (rect-center lrect) ehl)))
           (receive (es l-remain r-remain)
               (make-entry-string tw center (entry->string (entry tree)))
             (let ((left (expand-list (append ledges ls) h (make-string lw #\space)))
                   (right (expand-list (append redges rs) h (make-string rw #\space))))
               (let ((merged (map (lambda (l r)
                                    (append-space (string-append l " " r)
                                                  l-remain
                                                  r-remain))
                                  left
                                  right)))
                 (make-rect (cons es merged)
                            (+ tw l-remain r-remain)
                            (+ l-remain center)))))))))

(define (binary-tree->string tree)
  (apply string-append
         (map (cut string-append <> "\n")
              (rect-strings (binary-tree->rect tree)))))

(define (print-binary-tree tree)
  (print (binary-tree->string tree)))

(define (main . args)
  (until (read) eof-object? => x
    (print-binary-tree x)))

;; TEST
; (print-binary-tree '(A () B))
; (print-binary-tree '(1 30 (2 () ())))
; (print-binary-tree '(M (C () (E () ())) (S (R () ()) ())))
; (print-binary-tree '(1 () (2 () (3 () (4 () (5 () (6 () (7 8 9))))))))