github.com/jingcheng-WU/gonum@v0.9.1-0.20210323123734-f1a2a11a8f7b/stat/distmat/permutation.go (about)

     1  // Copyright ©2020 The Gonum 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  package distmat
     6  
     7  import (
     8  	"golang.org/x/exp/rand"
     9  
    10  	"github.com/jingcheng-WU/gonum/mat"
    11  )
    12  
    13  // UniformPermutation is a uniform distribution over the n!
    14  // permutation matrices of size n×n for a given n.
    15  type UniformPermutation struct {
    16  	rnd     *rand.Rand
    17  	indices []int
    18  }
    19  
    20  // NewUniformPermutation constructs a new permutation matrix
    21  // generator using the given random source.
    22  func NewUniformPermutation(src rand.Source) *UniformPermutation {
    23  	return &UniformPermutation{rnd: rand.New(src)}
    24  }
    25  
    26  // PermTo sets the given matrix to be a random permutation matrix.
    27  // It does not zero the destination's elements, so it is the responsibility
    28  // of the caller to ensure it is correctly conditioned prior to the call.
    29  //
    30  // PermTo panics if dst is not square.
    31  func (p *UniformPermutation) PermTo(dst *mat.Dense) {
    32  	r, c := dst.Dims()
    33  	if r != c {
    34  		panic(mat.ErrShape)
    35  	}
    36  	if r == 0 {
    37  		return
    38  	}
    39  	if len(p.indices) != r {
    40  		p.indices = make([]int, r)
    41  		for k := range p.indices {
    42  			p.indices[k] = k
    43  		}
    44  	}
    45  	p.rnd.Shuffle(r, func(i, j int) { p.indices[i], p.indices[j] = p.indices[j], p.indices[i] })
    46  	for i, j := range p.indices {
    47  		dst.Set(i, j, 1)
    48  	}
    49  }