diff --git a/src/main/java/kyu3/RailFenceCipher.java b/src/main/java/kyu3/RailFenceCipher.java index 4c5c22f..3de73e0 100644 --- a/src/main/java/kyu3/RailFenceCipher.java +++ b/src/main/java/kyu3/RailFenceCipher.java @@ -1,7 +1,6 @@ package kyu3; -import java.util.LinkedList; import java.util.Map; import java.util.TreeMap; @@ -40,56 +39,36 @@ static String decode(String s, int n) { // First computes the exact size of each rail, splits ciphertext into // contiguous rail segments, then reconstructs plaintext by replaying the // same rail traversal cycle. + if (n < 2) { + throw new IllegalArgumentException("Number of rails must be at least 2"); + } + if (s.isEmpty() || n >= s.length()) { + return s; + } + int[] size = new int[n]; Counter counter = new Counter(n - 1); for (int i = 0; i < s.length(); i++) { int tik = counter.tik(); size[tik]++; } - Map map = new TreeMap<>(); - StringBuilder result = new StringBuilder(); + int[] position = new int[n]; + int start = 0; for (int i = 0; i < n; i++) { - StringBuilder stringBuilder = new StringBuilder(); - if (i == 0) { - stringBuilder.append(s, 0, size[0]); - } else { - int start = sumArr(size, i); - int end = start + size[i]; - stringBuilder.append(s, start, end); - } - map.put(i, stringBuilder); - } - - Map> map2 = new TreeMap<>(); - - for (Map.Entry e : map.entrySet() - ) { - LinkedList strings = new LinkedList<>(); - for (char c : e.getValue().toString().toCharArray()) { - strings.add(String.valueOf(c)); - } - map2.put(e.getKey(), strings); + position[i] = start; + start += size[i]; } + StringBuilder result = new StringBuilder(s.length()); Counter counterRes = new Counter(n - 1); for (int i = 0; i < s.length(); i++) { int tik = counterRes.tik(); - result.append(map2.get(tik).poll()); + result.append(s.charAt(position[tik]++)); } return result.toString(); } - private static int sumArr(int[] arr, int i) { - // Sum widths of all previous rails to determine the start offset - // of the current rail segment. - int result = 0; - for (int j = 0; j < i; j++) { - result += arr[j]; - } - return result; - } - /** * Create two functions to encode and then decode a string using the Rail Fence Cipher. This cipher is used to encode * a string by placing each character successively in a diagonal along a set of "rails". First start off moving diagonally diff --git a/src/test/java/kyu3/RailFenceCipherTest.java b/src/test/java/kyu3/RailFenceCipherTest.java index ccc1fcc..41c73b7 100644 --- a/src/test/java/kyu3/RailFenceCipherTest.java +++ b/src/test/java/kyu3/RailFenceCipherTest.java @@ -20,6 +20,12 @@ void shouldEncodeAndDecodeKnownExample() { assertEquals("WEAREDISCOVEREDFLEEATONCE", RailFenceCipher.decode(encoded, 3)); } + @Test + void shouldDecodeWithoutAllocatingForUnusedRails() { + assertEquals("", RailFenceCipher.decode("", Integer.MAX_VALUE)); + assertEquals("abc", RailFenceCipher.decode("abc", Integer.MAX_VALUE)); + } + @ParameterizedTest @MethodSource("roundTripCases") void shouldRoundTripPlainTextForDifferentRails(String text, int rails) {