diff --git a/README.md b/README.md index 3e0739363..28d325260 100644 --- a/README.md +++ b/README.md @@ -244,6 +244,7 @@ All core data structures live in [`algorithms/data_structures/`](algorithms/data - [elias](algorithms/compression/elias.py) — Elias gamma and delta universal integer coding - [huffman_coding](algorithms/compression/huffman_coding.py) — variable-length prefix codes for lossless compression +- [lzw_compression](algorithms/compression/lzw_compression.py) — dictionary-based Lempel-Ziv-Welch compression - [rle_compression](algorithms/compression/rle_compression.py) — run-length encoding for consecutive character compression ### Dynamic Programming diff --git a/algorithms/compression/__init__.py b/algorithms/compression/__init__.py index 73b1a89cc..6571aec0c 100644 --- a/algorithms/compression/__init__.py +++ b/algorithms/compression/__init__.py @@ -1,5 +1,6 @@ from .elias import elias_delta, elias_gamma from .huffman_coding import HuffmanCoding +from .lzw_compression import lzw_decode, lzw_encode from .rle_compression import decode_rle, encode_rle __all__ = [ @@ -8,4 +9,6 @@ "elias_delta", "elias_gamma", "encode_rle", + "lzw_decode", + "lzw_encode", ] diff --git a/algorithms/compression/lzw_compression.py b/algorithms/compression/lzw_compression.py new file mode 100644 index 000000000..01d7b38c6 --- /dev/null +++ b/algorithms/compression/lzw_compression.py @@ -0,0 +1,120 @@ +""" +Lempel-Ziv-Welch (LZW) Compression + +A dictionary-based lossless compression algorithm. It builds a dictionary of +substrings during encoding and replaces repeated substrings with dictionary +codes. Decompression reconstructs the same dictionary on the fly to recover +the original data. + +Reference: https://en.wikipedia.org/wiki/Lempel%E2%80%93Ziv%E2%80%93Welch + +Complexity: + Time: O(n) for both encoding and decoding + Space: O(n) for the dictionary/code table +""" + +from __future__ import annotations + + +def _build_initial_dictionary(data: str) -> dict[str, int]: + """Create the initial dictionary containing all unique input characters. + + Args: + data: The input string. + + Returns: + A dictionary mapping characters to integer codes. + """ + return {char: index for index, char in enumerate(sorted(set(data)))} + + +def lzw_encode(data: str) -> tuple[list[int], dict[int, str]]: + """Compress a string using the LZW algorithm. + + Args: + data: The input string to compress. + + Returns: + A tuple of the integer codes representing the compressed data and the + initial code-to-character dictionary needed for decoding. + + Examples: + >>> codes, dictionary = lzw_encode("ABABABA") + >>> codes + [0, 1, 2, 4] + >>> dictionary + {0: 'A', 1: 'B'} + >>> lzw_encode("") + ([], {}) + """ + if not data: + return [], {} + + dictionary = _build_initial_dictionary(data) + next_code = len(dictionary) + encoded: list[int] = [] + current: str = "" + + for char in data: + combined = current + char + if combined in dictionary: + current = combined + else: + encoded.append(dictionary[current]) + dictionary[combined] = next_code + next_code += 1 + current = char + + if current: + encoded.append(dictionary[current]) + + initial_dictionary = { + code: char + for char, code in dictionary.items() + if len(char) == 1 + } + return encoded, initial_dictionary + + +def lzw_decode(encoded: list[int], initial_dictionary: dict[int, str]) -> str: + """Decompress a list of LZW codes back into the original string. + + Args: + encoded: The list of integer codes produced by lzw_encode. + initial_dictionary: Mapping of initial codes to single-character + strings, as returned by lzw_encode. + + Returns: + The decoded original string. + + Examples: + >>> lzw_decode([0, 1, 2, 4], {0: "A", 1: "B"}) + 'ABABABA' + >>> lzw_decode([], {}) + '' + """ + if not encoded: + return "" + + codes_to_strings = dict(initial_dictionary) + next_code = max(codes_to_strings.keys()) + 1 + decoded: str = "" + previous: str = "" + + for code in encoded: + if code in codes_to_strings: + current = codes_to_strings[code] + elif code == next_code and previous: + current = previous + previous[0] + else: + raise ValueError(f"Invalid LZW code: {code}") + + decoded += current + + if previous: + codes_to_strings[next_code] = previous + current[0] + next_code += 1 + + previous = current + + return decoded diff --git a/tests/test_compression.py b/tests/test_compression.py index ca9156abc..80b183253 100644 --- a/tests/test_compression.py +++ b/tests/test_compression.py @@ -2,6 +2,7 @@ from algorithms.compression.elias import elias_delta, elias_gamma from algorithms.compression.huffman_coding import HuffmanCoding +from algorithms.compression.lzw_compression import lzw_decode, lzw_encode from algorithms.compression.rle_compression import decode_rle, encode_rle @@ -101,5 +102,25 @@ def test_elias_delta(self): self.assertEqual(correct_result, result) +class TestLZWCompression(unittest.TestCase): + def test_lzw_encode(self): + codes, dictionary = lzw_encode("ABABABA") + self.assertEqual([0, 1, 2, 4], codes) + self.assertEqual({0: "A", 1: "B"}, dictionary) + + def test_lzw_decode(self): + self.assertEqual("ABABABA", lzw_decode([0, 1, 2, 4], {0: "A", 1: "B"})) + + def test_lzw_roundtrip(self): + data = "TOBEORNOTTOBEORTOBEORNOT" + encoded, dictionary = lzw_encode(data) + decoded = lzw_decode(encoded, dictionary) + self.assertEqual(data, decoded) + + def test_lzw_empty(self): + self.assertEqual(([], {}), lzw_encode("")) + self.assertEqual("", lzw_decode([], {})) + + if __name__ == "__main__": unittest.main()