Submission #5421738

---

Source Code Expand

(eval-when (:compile-toplevel :load-toplevel :execute)
  (defparameter OPT
    #+swank '(optimize (speed 3) (safety 2))
    #-swank '(optimize (speed 3) (safety 0) (debug 0)))
  #+swank (progn (ql:quickload '(:cl-debug-print :fiveam))
                 (shadow :run)
                 (use-package :fiveam)))
#+swank (cl-syntax:use-syntax cl-debug-print:debug-print-syntax)

;; BEGIN_INSERTED_CONTENTS
(declaim (type (integer 0 #.most-positive-fixnum) *recursion-depth*))
(defparameter *recursion-depth* 0)

(defmacro with-cache (cache-attribs def-form)
  (let* ((cache-attribs (if (atom cache-attribs) (list cache-attribs) cache-attribs))
         (cache-type (first cache-attribs))
         (dimensions-with-* (when (eql cache-type :array) (second cache-attribs)))
         (dimensions (remove '* dimensions-with-*))
         (rank (length dimensions))
         (rest-attribs (ecase cache-type
                         (:hash-table (cdr cache-attribs))
                         (:array (cddr cache-attribs))))
         (key (prog1 (getf rest-attribs :key) (remf rest-attribs :key)))
         (debug (prog1 (getf rest-attribs :debug) (remf rest-attribs :debug)))
         (cache-form (case cache-type
                       (:hash-table `(make-hash-table ,@rest-attribs))
                       (:array `(make-array (list ,@dimensions) ,@rest-attribs))))
         (initial-element (when (eql cache-type :array)
                            (assert (member :initial-element rest-attribs))
                            (getf rest-attribs :initial-element))))
    (let ((cache (gensym))
          (value (gensym))
	  (present-p (gensym))
          (name-alias (gensym))
	  (args-lst (gensym))
          (indices (loop repeat rank collect (gensym))))
      (labels ((debug (name args obj)
                 (let ((value (gensym)))
                   (if debug
                       `(progn
                          (format t "~A~A: (~A ~{~A~^ ~}) =>~%"
                                  (make-string *recursion-depth*
                                               :element-type 'base-char
                                               :initial-element #\ )
                                  *recursion-depth*
                                  ',name
                                  (list ,@args))
                          (let ((,value (let ((*recursion-depth* (1+ *recursion-depth*)))
                                          ,obj)))
                            (format t "~A~A: (~A ~{~A~^ ~}) => ~A~%"
                                    (make-string *recursion-depth*
                                               :element-type 'base-char
                                               :initial-element #\ )
                                    *recursion-depth*
                                    ',name
                                    (list ,@args)
                                    ,value)
                            ,value))
                       obj)))
               (make-cache-check-form (cache-type name args)
                 (debug name
                        args
                        (case cache-type
                          (:hash-table
                           `(let ((,args-lst (funcall ,(or key #'list) ,@args)))
                              (multiple-value-bind (,value ,present-p)
                                  (gethash ,args-lst ,cache)
                                (if ,present-p
                                    ,value
                                    (setf (gethash ,args-lst ,cache)
                                          (,name-alias ,@args))))))
                          (:array
                           (let ((memoized-args (loop for dimension in dimensions-with-*
                                                      for arg in args
                                                      unless (eql dimension '*)
                                                      collect arg)))
                             (if key
                                 `(multiple-value-bind ,indices
                                      (funcall ,key ,@memoized-args)
                                    (let ((,value (aref ,cache ,@indices)))
                                      (if (eql ,initial-element ,value)
                                          (setf (aref ,cache ,@indices)
                                                (,name-alias ,@args))
                                          ,value)))
                                 `(let ((,value (aref ,cache ,@memoized-args)))
                                    (if (eql ,initial-element ,value)
                                        (setf (aref ,cache ,@memoized-args)
                                              (,name-alias ,@args))
                                        ,value))))))))
               (make-reset-form (cache-type)
                 (case cache-type
                   (:hash-table `(setf ,cache (make-hash-table ,@rest-attribs)))
                   (:array `(prog1 nil
                              (fill (array-storage-vector ,cache) ,initial-element)))))
               (make-reset-name (name)
                 (intern (format nil "RESET-~A" (symbol-name name))))
               (extract-declarations (body)
                 (remove-if-not (lambda (form) (eql 'declare (car form))) body)))
        (ecase (car def-form)
          ((defun)
           (destructuring-bind (_ name args &body body) def-form
             (declare (ignore _))
             `(let ((,cache ,cache-form))
                (defun ,(make-reset-name name) () ,(make-reset-form cache-type))
                (defun ,name ,args
                  ,@(extract-declarations body)
                  (labels ((,name-alias ,args ,@body))
                    (declare (inline ,name-alias))
                    ,(make-cache-check-form cache-type name args))))))
          ((nlet)
           (destructuring-bind (_ name bindings &body body) def-form
             (declare (ignore _))
             `(let ((,cache ,cache-form))
                (nlet ,name ,bindings
                  ,@(extract-declarations body)
                  ,(let ((args (mapcar (lambda (x) (if (atom x) x (car x))) bindings)))
                     `(labels ((,name-alias ,args ,@body))
                        (declare (inline ,name-alias))
                        ,(make-cache-check-form cache-type name args)))))))
          ((labels flet)
           (destructuring-bind (_ definitions &body labels-body) def-form
             (declare (ignore _))
             (destructuring-bind (name args &body body) (car definitions)
               `(let ((,cache ,cache-form))
                  (,(car def-form)
                   ((,(make-reset-name name) () ,(make-reset-form cache-type))
                    (,name ,args
                           ,@(extract-declarations body)
                           (labels ((,name-alias ,args ,@body))
                             (declare (inline ,name-alias))
                             ,(make-cache-check-form cache-type name args)))
                    ,@(cdr definitions))
                   (declare (ignorable #',(make-reset-name name)))
                   ,@labels-body))))))))))

(declaim (ftype (function * (values fixnum &optional)) read-fixnum))
(defun read-fixnum (&optional (in *standard-input*))
  (declare #.OPT)
  (macrolet ((%read-byte ()
               `(the (unsigned-byte 8)
                     #+swank (char-code (read-char in nil #\Nul))
                     #-swank (sb-impl::ansi-stream-read-byte in nil #.(char-code #\Nul) nil))))
    (let* ((minus nil)
           (result (loop (let ((byte (%read-byte)))
                           (cond ((<= 48 byte 57)
                                  (return (- byte 48)))
                                 ((zerop byte) ; #\Nul
                                  ;; (return-from read-fixnum 0)
                                  (error "Read EOF or #\Nul."))
                                 ((= byte #.(char-code #\-))
                                  (setf minus t)))))))
      (declare ((integer 0 #.most-positive-fixnum) result))
      (loop
        (let* ((byte (%read-byte)))
          (if (<= 48 byte 57)
              (setq result (+ (- byte 48) (* 10 (the (integer 0 #.(floor most-positive-fixnum 10)) result))))
              (return (if minus (- result) result))))))))

;;;
;;; Binomial coefficient with mod
;;; build: O(n)
;;; query: O(1)
;;;

(defconstant +binom-size+ 2000)
(defconstant +binom-mod+ #.(+ (expt 10 9) 7))

(declaim ((simple-array (unsigned-byte 32) (*)) *fact* *fact-inv* *inv*))
(defparameter *fact* (make-array +binom-size+ :element-type '(unsigned-byte 32)))
(defparameter *fact-inv* (make-array +binom-size+ :element-type '(unsigned-byte 32)))
(defparameter *inv* (make-array +binom-size+ :element-type '(unsigned-byte 32)))

(defun initialize-binom ()
  (setf (aref *fact* 0) 1
        (aref *fact* 1) 1
        (aref *fact-inv* 0) 1
        (aref *fact-inv* 1) 1
        (aref *inv* 1) 1)
  (loop for i from 2 below +binom-size+
        do (setf (aref *fact* i) (mod (* i (aref *fact* (- i 1))) +binom-mod+)
                 (aref *inv* i) (mod (- (* (aref *inv* (rem +binom-mod+ i))
                                           (floor +binom-mod+ i)))
                                     +binom-mod+)
                 (aref *fact-inv* i) (mod (* (aref *inv* i)
                                             (aref *fact-inv* (- i 1)))
                                          +binom-mod+))))

(initialize-binom)

(declaim (inline binom))
(defun binom (n k)
  (if (or (< n k) (< n 0) (< k 0))
      0
      (mod (* (aref *fact* n)
              (mod (* (aref *fact-inv* k) (aref *fact-inv* (- n k))) +binom-mod+))
           +binom-mod+)))

(declaim (inline multinomial))
(defun multinomial (&rest ks)
  (let ((sum 0)
        (result 1))
    (declare ((integer 0 #.most-positive-fixnum) result sum))
    (dolist (k ks)
      (incf sum (the fixnum k))
      (setq result
            (mod (* result (aref *fact-inv* k)) +binom-mod+)))
    (mod (* result (aref *fact* sum)) +binom-mod+)))

(defmacro dbg (&rest forms)
  #+swank
  (if (= (length forms) 1)
      `(format *error-output* "~A => ~A~%" ',(car forms) ,(car forms))
      `(format *error-output* "~A => ~A~%" ',forms `(,,@forms)))
  #-swank (declare (ignore forms)))

(defmacro define-int-types (&rest bits)
  `(progn
     ,@(mapcar (lambda (b) `(deftype ,(intern (format nil "UINT~A" b)) () '(unsigned-byte ,b))) bits)
     ,@(mapcar (lambda (b) `(deftype ,(intern (format nil "INT~A" b)) () '(signed-byte ,b))) bits)))
(define-int-types 2 4 7 8 15 16 31 32 62 63 64)

(declaim (inline println))
(defun println (obj &optional (stream *standard-output*))
  (let ((*read-default-float-format* 'double-float))
    (prog1 (princ obj stream) (terpri stream))))

(defconstant +mod+ 1000000007)

(defun mod* (&rest args)
  (reduce (lambda (x y) (mod (* x y) +mod+)) args))

(define-compiler-macro mod* (&rest args)
  (if (null args)
      1
      (reduce (lambda (x y) `(mod (* ,x ,y) +mod+)) args)))

(defun mod+ (&rest args)
  (reduce (lambda (x y) (mod (+ x y) +mod+)) args))

(define-compiler-macro mod+ (&rest args)
  (if (null args)
      0
      (reduce (lambda (x y) `(mod (+ ,x ,y) +mod+)) args)))

(define-modify-macro incfmod (delta divisor)
  (lambda (x y divisor) (mod (+ x y) divisor)))

;; Body

(defun main ()
  (declare #.OPT)
  (let* ((n (read))
         (graph (make-array n :element-type 'list :initial-element nil))
         (as (make-array (- n 1) :element-type 'uint32))
         (bs (make-array (- n 1) :element-type 'uint32)))
    (with-cache (:hash-table :test #'eql
                             :key #'(lambda (x y)
                                      (the uint62 (+ (ash x 31) y)))
                             :size n)
      (labels ((subtree-size (parent top)
                 (declare (uint32 parent top)
                          (values uint32))
                 (+ 1
                    (loop for child of-type uint32 in (aref graph top)
                          unless (= child parent)
                          sum (subtree-size top child) of-type uint32))))
        (with-cache (:hash-table :test #'eql
                                 :key #'(lambda (x y)
                                          (the uint62 (+ (ash x 31) y)))
                                 :size n)
          (labels ((subtree-number (parent top)
                     (declare (uint32 parent top)
                              (values uint32))
                     (let ((res 1)
                           args)
                       (declare (uint32 res))
                       (dolist (neighbor (aref graph top))
                         (declare (uint32 neighbor))
                         (unless (= neighbor parent)
                           (setf res (mod* res (subtree-number top neighbor)))
                           (push (subtree-size top neighbor) args)))
                       (mod* res (apply #'multinomial args)))))
            (dotimes (i (- n 1))
              (let ((a (- (read-fixnum) 1))
                    (b (- (read-fixnum) 1)))
                (push a (aref graph b))
                (push b (aref graph a))
                (setf (aref as i) a)
                (setf (aref bs i) b)))
            (let ((res 0))
              (declare (uint32 res))
              (loop for a across as
                    for b across bs
                    do (incfmod res
                                (mod* (multinomial (- (subtree-size a b) 1)
                                                   (- (subtree-size b a) 1))
                                      (subtree-number a b)
                                      (subtree-number b a))
                                +mod+))
              (println res))))))))

#-swank(main)

Submission Info

Judge Result