Home | History | Annotate | Download | only in codewalk 
      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 }
    131