1 files changed, 12 insertions, 6 deletions
diff --git a/ass2/q2/suffix_array.py b/ass2/q2/suffix_array.py
index e743180..edfaa80 100644
--- a/ass2/q2/suffix_array.py
+++ b/ass2/q2/suffix_array.py
@@ -15,18 +15,24 @@ def read_in_string(filename):
         return file.read()
 
 
-def write_suffix_array(tree, filename):
+def suffix_array(string):
+    tree = ukkonen(string)
+    buffer = []
+    depth_first_apply(tree, buffer.append)
+    return buffer
+
+
+def write_suffix_array(string, filename):
     with open(filename, "w") as file:
-        buffer = []
-        depth_first_apply(tree, buffer.append)
+        buffer = suffix_array(string)
         file.write("\n".join(map(str, buffer)))
 
 
 def main():
     assert len(sys.argv) == 2
     string = read_in_string(sys.argv[1])
-    tree = ukkonen(string)
-    write_suffix_array(tree, "output_suffix_array.txt")
+    write_suffix_array(string, "output_suffix_array.txt")
 
 
-main()
+if __name__ == "__main__":
+    main()