<commit_before>package bwt

import (
	"errors"
	"sort"

	"github.com/shenwei356/util/byteutil"
)

// ErrEndSymbolExisted means you should choose another EndSymbol
var ErrEndSymbolExisted = errors.New("bwt: end-symbol existed in string")

// Transform returns Burrows–Wheeler transform  of a byte slice.
// See https://en.wikipedia.org/wiki/Burrows%E2%80%93Wheeler_transform
func Transform(s []byte, es byte) ([]byte, [][]byte, error) {
	count := byteutil.CountOfByte(s)
	if _, ok := count[es]; ok {
		return nil, nil, ErrEndSymbolExisted
	}
	s = append(s, es)
	n := len(s)

	rotations := make([][]byte, n)
	i := 0
	for j := 0; j < n; j++ {
		rotations[i] = append(s[n-j:], s[0:n-j]...)
		i++
	}
	sort.Sort(byteutil.SliceOfByteSlice(rotations))

	bwt := make([]byte, n)
	i = 0
	for _, t := range rotations {
		bwt[i] = t[n-1]
		i++
	}
	return bwt, rotations, nil
}

// InverseTransform reverses the bwt to original byte slice
func InverseTransform(t []byte, es byte) []byte {
	n := len(t)
	lines := make([][]byte, n)
	for i := 0; i < n; i++ {
		lines[i] = make([]byte, n)
	}

	for i := 0; i < n; i++ {
		for j := 0; j < n; j++ {
			lines[j][n-1-i] = t[j]
		}
		sort.Sort(byteutil.SliceOfByteSlice(lines))
	}

	s := make([]byte, n-1)
	for _, line := range lines {
		if line[n-1] == es {
			s = line[0 : n-1]
			break
		}
	}
	return s
}

// SuffixArray returns the suffix array of s
func SuffixArray(s []byte) []int {
	n := len(s)
	suffixMap := make(map[string]int, n)
	for i := 0; i < n; i++ {
		suffixMap[string(s[i:])] = i
	}
	suffixes := make([]string, n)
	i := 0
	for suffix := range suffixMap {
		suffixes[i] = suffix
		i++
	}
	indice := make([]int, n+1)
	indice[0] = n
	i = 1
	sort.Strings(suffixes)
	for _, suffix := range suffixes {
		indice[i] = suffixMap[suffix]
		i++
	}
	return indice
}

// ErrInvalidSuffixArray means length of sa is not equal to 1+len(s)
var ErrInvalidSuffixArray = errors.New("bwt: invalid suffix array")

// FromSuffixArray compute BWT from sa
func FromSuffixArray(s []byte, sa []int, es byte) ([]byte, error) {
	if len(s)+1 != len(sa) || sa[0] != len(s) {
		return nil, ErrInvalidSuffixArray
	}
	bwt := make([]byte, len(sa))
	bwt[0] = s[len(s)-1]
	for i := 1; i < len(sa); i++ {
		if sa[i] == 0 {
			bwt[i] = es
		} else {
			bwt[i] = s[sa[i]-1]
		}
	}
	return bwt, nil
}<commit_msg>optimize SuffixArray<commit_after>package bwt

import (
	"bytes"
	"errors"
	"sort"

	"github.com/shenwei356/util/byteutil"
)

// ErrEndSymbolExisted means you should choose another EndSymbol
var ErrEndSymbolExisted = errors.New("bwt: end-symbol existed in string")

// Transform returns Burrows–Wheeler transform  of a byte slice.
// See https://en.wikipedia.org/wiki/Burrows%E2%80%93Wheeler_transform
func Transform(s []byte, es byte) ([]byte, [][]byte, error) {
	count := byteutil.CountOfByte(s)
	if _, ok := count[es]; ok {
		return nil, nil, ErrEndSymbolExisted
	}
	s = append(s, es)
	n := len(s)

	rotations := make([][]byte, n)
	i := 0
	for j := 0; j < n; j++ {
		rotations[i] = append(s[n-j:], s[0:n-j]...)
		i++
	}
	sort.Sort(byteutil.SliceOfByteSlice(rotations))

	bwt := make([]byte, n)
	i = 0
	for _, t := range rotations {
		bwt[i] = t[n-1]
		i++
	}
	return bwt, rotations, nil
}

// InverseTransform reverses the bwt to original byte slice
func InverseTransform(t []byte, es byte) []byte {
	n := len(t)
	lines := make([][]byte, n)
	for i := 0; i < n; i++ {
		lines[i] = make([]byte, n)
	}

	for i := 0; i < n; i++ {
		for j := 0; j < n; j++ {
			lines[j][n-1-i] = t[j]
		}
		sort.Sort(byteutil.SliceOfByteSlice(lines))
	}

	s := make([]byte, n-1)
	for _, line := range lines {
		if line[n-1] == es {
			s = line[0 : n-1]
			break
		}
	}
	return s
}

// SuffixArray returns the suffix array of s
func SuffixArray(s []byte) []int {
	sa := make([]int, len(s)+1)
	sa[0] = len(s)

	for i := 0; i < len(s); i++ {
		sa[i+1] = i
	}
	sort.Slice(sa[1:], func(i, j int) bool {
		return bytes.Compare(s[sa[i+1]:], s[sa[j+1]:]) < 0
	})
	return sa
}

// ErrInvalidSuffixArray means length of sa is not equal to 1+len(s)
var ErrInvalidSuffixArray = errors.New("bwt: invalid suffix array")

// FromSuffixArray compute BWT from sa
func FromSuffixArray(s []byte, sa []int, es byte) ([]byte, error) {
	if len(s)+1 != len(sa) || sa[0] != len(s) {
		return nil, ErrInvalidSuffixArray
	}
	bwt := make([]byte, len(sa))
	bwt[0] = s[len(s)-1]
	for i := 1; i < len(sa); i++ {
		if sa[i] == 0 {
			bwt[i] = es
		} else {
			bwt[i] = s[sa[i]-1]
		}
	}
	return bwt, nil
}
{'fraction_non_alphanumeric': 0.18791946308724833, 'fraction_numerical': 0.015883668903803133, 'mean_word_length': 3.4049261083743843, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '837709', 'n_tokens_mistral': 1984, 'n_tokens_neox': 1865, 'n_words': 498}