github.com/mdempsky/go@v0.0.0-20151201204031-5dd372bd1e70/doc/codewalk/markov.go (about)

     1  // Copyright 2011 The Go Authors.  All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  /*
     6  Generating random text: a Markov chain algorithm
     7  
     8  Based on the program presented in the "Design and Implementation" chapter
     9  of The Practice of Programming (Kernighan and Pike, Addison-Wesley 1999).
    10  See also Computer Recreations, Scientific American 260, 122 - 125 (1989).
    11  
    12  A Markov chain algorithm generates text by creating a statistical model of
    13  potential textual suffixes for a given prefix. Consider this text:
    14  
    15  	I am not a number! I am a free man!
    16  
    17  Our Markov chain algorithm would arrange this text into this set of prefixes
    18  and suffixes, or "chain": (This table assumes a prefix length of two words.)
    19  
    20  	Prefix       Suffix
    21  
    22  	"" ""        I
    23  	"" I         am
    24  	I am         a
    25  	I am         not
    26  	a free       man!
    27  	am a         free
    28  	am not       a
    29  	a number!    I
    30  	number! I    am
    31  	not a        number!
    32  
    33  To generate text using this table we select an initial prefix ("I am", for
    34  example), choose one of the suffixes associated with that prefix at random
    35  with probability determined by the input statistics ("a"),
    36  and then create a new prefix by removing the first word from the prefix
    37  and appending the suffix (making the new prefix is "am a"). Repeat this process
    38  until we can't find any suffixes for the current prefix or we exceed the word
    39  limit. (The word limit is necessary as the chain table may contain cycles.)
    40  
    41  Our version of this program reads text from standard input, parsing it into a
    42  Markov chain, and writes generated text to standard output.
    43  The prefix and output lengths can be specified using the -prefix and -words
    44  flags on the command-line.
    45  */
    46  package main
    47  
    48  import (
    49  	"bufio"
    50  	"flag"
    51  	"fmt"
    52  	"io"
    53  	"math/rand"
    54  	"os"
    55  	"strings"
    56  	"time"
    57  )
    58  
    59  // Prefix is a Markov chain prefix of one or more words.
    60  type Prefix []string
    61  
    62  // String returns the Prefix as a string (for use as a map key).
    63  func (p Prefix) String() string {
    64  	return strings.Join(p, " ")
    65  }
    66  
    67  // Shift removes the first word from the Prefix and appends the given word.
    68  func (p Prefix) Shift(word string) {
    69  	copy(p, p[1:])
    70  	p[len(p)-1] = word
    71  }
    72  
    73  // Chain contains a map ("chain") of prefixes to a list of suffixes.
    74  // A prefix is a string of prefixLen words joined with spaces.
    75  // A suffix is a single word. A prefix can have multiple suffixes.
    76  type Chain struct {
    77  	chain     map[string][]string
    78  	prefixLen int
    79  }
    80  
    81  // NewChain returns a new Chain with prefixes of prefixLen words.
    82  func NewChain(prefixLen int) *Chain {
    83  	return &Chain{make(map[string][]string), prefixLen}
    84  }
    85  
    86  // Build reads text from the provided Reader and
    87  // parses it into prefixes and suffixes that are stored in Chain.
    88  func (c *Chain) Build(r io.Reader) {
    89  	br := bufio.NewReader(r)
    90  	p := make(Prefix, c.prefixLen)
    91  	for {
    92  		var s string
    93  		if _, err := fmt.Fscan(br, &s); err != nil {
    94  			break
    95  		}
    96  		key := p.String()
    97  		c.chain[key] = append(c.chain[key], s)
    98  		p.Shift(s)
    99  	}
   100  }
   101  
   102  // Generate returns a string of at most n words generated from Chain.
   103  func (c *Chain) Generate(n int) string {
   104  	p := make(Prefix, c.prefixLen)
   105  	var words []string
   106  	for i := 0; i < n; i++ {
   107  		choices := c.chain[p.String()]
   108  		if len(choices) == 0 {
   109  			break
   110  		}
   111  		next := choices[rand.Intn(len(choices))]
   112  		words = append(words, next)
   113  		p.Shift(next)
   114  	}
   115  	return strings.Join(words, " ")
   116  }
   117  
   118  func main() {
   119  	// Register command-line flags.
   120  	numWords := flag.Int("words", 100, "maximum number of words to print")
   121  	prefixLen := flag.Int("prefix", 2, "prefix length in words")
   122  
   123  	flag.Parse()                     // Parse command-line flags.
   124  	rand.Seed(time.Now().UnixNano()) // Seed the random number generator.
   125  
   126  	c := NewChain(*prefixLen)     // Initialize a new Chain.
   127  	c.Build(os.Stdin)             // Build chains from standard input.
   128  	text := c.Generate(*numWords) // Generate text.
   129  	fmt.Println(text)             // Write text to standard output.
   130  }