Add LZW compression module

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

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",
 ]

algorithms/compression/lzw_compression.py

@@ -0,0 +1,116 @@
+"""
+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

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()