1 files changed, 9 insertions, 8 deletions

diff --git a/srchr/src/corpus.rs b/srchr/src/corpus.rs
index 7216b1a..306018c 100644
--- a/srchr/src/corpus.rs
+++ b/srchr/src/corpus.rs
@@ -19,7 +19,7 @@ const NUM_BIGRAMS: usize = NUM_KEYS * NUM_KEYS;
 pub struct Corpus {
     bigram_percs: [f32; NUM_BIGRAMS],
     skipgram_percs: [f32; NUM_BIGRAMS],
-    character_count: [u32; 128],
+    character_percs: [f32; 128],
     total_count: u32,
     total_bigrams: u32,
 }
@@ -31,8 +31,8 @@ fn pair_to_index(x: u8, y: u8) -> usize {
 }
 
 impl Corpus {
-    pub fn get_character_count(&self, c: u8) -> u32 {
-        self.character_count[c as usize]
+    pub fn get_character_perc(&self, c: u8) -> f32 {
+        self.character_percs[c as usize]
     }
 
     pub fn get_bigram_perc(&self, x: u8, y: u8) -> f32 {
@@ -43,10 +43,6 @@ impl Corpus {
         self.skipgram_percs[pair_to_index(x, y)]
     }
 
-    pub fn get_total_count(&self) -> u32 {
-        self.total_count
-    }
-
     pub fn get_total_bigrams(&self) -> u32 {
         self.total_bigrams
     }
@@ -89,16 +85,21 @@ impl Corpus {
 
         let mut bigram_percs = [0.0; NUM_BIGRAMS];
         let mut skipgram_percs = [0.0; NUM_BIGRAMS];
+        let mut character_percs = [0.0; 128];
 
         for i in 0..NUM_BIGRAMS {
             bigram_percs[i] = bigram_count[i] as f32 / total_bigrams as f32;
             skipgram_percs[i] = skipgram_count[i] as f32 / total_trigrams as f32;
         }
 
+        for i in 0..128 {
+            character_percs[i] = character_count[i] as f32 / total_count as f32;
+        }
+
         return Ok(Corpus {
             skipgram_percs,
             bigram_percs,
-            character_count,
+            character_percs,
             total_count,
             total_bigrams,
         });