Submission #5421620


Source Code Expand

;; -*- coding: utf-8 -*-
(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
;; (with-cache (:hash-table :test #'equal :key #'cons)
;;   (defun ...))
;; (with-cache (:array (10 10 * 10) :initial-element -1 :element-type 'fixnum)
;;   (defun foo (a b c d) ...) ; C is ignored.
(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))))))))))

;; (test with-cache
;;   (finishes (macroexpand `(with-cache (:hash-table :test #'equal)
;;                             (defun add (x y) (+ x y)))))
;;   (finishes (macroexpand `(with-cache (:array '(10 10)
;;                                            :element-type 'fixnum
;;                                            :initial-element -1)
;;                             (defun add (x y) (+ x y)))))
;;   (finishes (macroexpand `(with-cache (:array '(10 10)
;;                                            :element-type 'fixnum
;;                                            :initial-element -1)
;;                             (labels ((add (x y) (+ x y))
;; 		                     (my-print (x) (print x)))
;; 	                      (add 1 2))))))

(declaim (ftype (function * (values fixnum &optional)) read-fixnum))
(defun read-fixnum (&optional (in *standard-input*))
  (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+ 510000)
(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 ()
  (let* ((n (read))
         (graph (make-array n :element-type 'list :initial-element nil))
         edges)
    (with-cache (:array (1000 1000) :initial-element #xffff :element-type 'uint16)
      (labels ((subtree-size (parent top)
                 (declare (uint32 parent top))
                 (+ 1
                    (loop for child of-type uint32 in (aref graph top)
                          unless (= child parent)
                          sum (subtree-size top child) of-type uint32))))
        (with-cache (:array (1000 1000)
                     :initial-element #xffffffff
                     :element-type 'uint32)
          (labels ((subtree-number (parent top)
                     (declare (uint32 parent top))
                     (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))
                (push (cons a b) edges)))
            (let ((res 0))
              (declare (uint32 res))
              (loop for (a . b) in edges
                    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

Submission Time
Task N - 木
User sansaqua
Language Common Lisp (SBCL 1.1.14)
Score 5
Code Size 14621 Byte
Status AC
Exec Time 255 ms
Memory 63968 KB

Judge Result

Set Name All
Score / Max Score 5 / 5
Status
AC × 9
Set Name Test Cases
All 00, 01, 02, 03, 04, 05, 06, 90, 91
Case Name Status Exec Time Memory
00 AC 241 ms 47204 KB
01 AC 255 ms 63968 KB
02 AC 190 ms 43488 KB
03 AC 191 ms 43492 KB
04 AC 193 ms 43492 KB
05 AC 191 ms 43492 KB
06 AC 191 ms 43492 KB
90 AC 191 ms 43488 KB
91 AC 190 ms 43492 KB