{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a: 523 +-0++0+\n",
      "b: 10 +0+\n",
      "c: 65 +-++-\n",
      "a * (b - c): -28765 ---0--0-0-\n"
     ]
    }
   ],
   "source": [
    "class BalancedTernary:\n",
    "    # Represented as a list of 0, 1 or -1s, with least significant digit first.\n",
    " \n",
    "    str2dig = {'+': 1, '-': -1, '0': 0} # immutable\n",
    "    dig2str = {1: '+', -1: '-', 0: '0'} # immutable\n",
    "    table = ((0, -1), (1, -1), (-1, 0), (0, 0), (1, 0), (-1, 1), (0, 1)) # immutable\n",
    " \n",
    "    def __init__(self, inp):\n",
    "        if isinstance(inp, str):\n",
    "            self.digits = [BalancedTernary.str2dig[c] for c in reversed(inp)]\n",
    "        elif isinstance(inp, int):\n",
    "            self.digits = self._int2ternary(inp)\n",
    "        elif isinstance(inp, BalancedTernary):\n",
    "            self.digits = list(inp.digits)\n",
    "        elif isinstance(inp, list):\n",
    "            if all(d in (0, 1, -1) for d in inp):\n",
    "                self.digits = list(inp)\n",
    "            else:\n",
    "                raise ValueError(\"BalancedTernary: Wrong input digits.\")\n",
    "        else:\n",
    "            raise TypeError(\"BalancedTernary: Wrong constructor input.\")\n",
    " \n",
    "    @staticmethod\n",
    "    def _int2ternary(n):\n",
    "        if n == 0: return []\n",
    "        if (n % 3) == 0: return [0] + BalancedTernary._int2ternary(n // 3)\n",
    "        if (n % 3) == 1: return [1] + BalancedTernary._int2ternary(n // 3)\n",
    "        if (n % 3) == 2: return [-1] + BalancedTernary._int2ternary((n + 1) // 3)\n",
    " \n",
    "    def to_int(self):\n",
    "        return reduce(lambda y,x: x + 3 * y, reversed(self.digits), 0)\n",
    " \n",
    "    def __repr__(self):\n",
    "        if not self.digits: return \"0\"\n",
    "        return \"\".join(BalancedTernary.dig2str[d] for d in reversed(self.digits))\n",
    " \n",
    "    @staticmethod\n",
    "    def _neg(digs):\n",
    "        return [-d for d in digs]\n",
    " \n",
    "    def __neg__(self):\n",
    "        return BalancedTernary(BalancedTernary._neg(self.digits))\n",
    " \n",
    "    @staticmethod\n",
    "    def _add(a, b, c=0):\n",
    "        if not (a and b):\n",
    "            if c == 0:\n",
    "                return a or b\n",
    "            else:\n",
    "                return BalancedTernary._add([c], a or b)\n",
    "        else:\n",
    "            (d, c) = BalancedTernary.table[3 + (a[0] if a else 0) + (b[0] if b else 0) + c]\n",
    "            res = BalancedTernary._add(a[1:], b[1:], c)\n",
    "            # trim leading zeros\n",
    "            if res or d != 0:\n",
    "                return [d] + res\n",
    "            else:\n",
    "                return res\n",
    " \n",
    "    def __add__(self, b):\n",
    "        return BalancedTernary(BalancedTernary._add(self.digits, b.digits))\n",
    " \n",
    "    def __sub__(self, b):\n",
    "        return self + (-b)\n",
    " \n",
    "    @staticmethod\n",
    "    def _mul(a, b):\n",
    "        if not (a and b):\n",
    "            return []\n",
    "        else:\n",
    "            if   a[0] == -1: x = BalancedTernary._neg(b)\n",
    "            elif a[0] ==  0: x = []\n",
    "            elif a[0] ==  1: x = b\n",
    "            else: assert False\n",
    "            y = [0] + BalancedTernary._mul(a[1:], b)\n",
    "            return BalancedTernary._add(x, y)\n",
    " \n",
    "    def __mul__(self, b):\n",
    "        return BalancedTernary(BalancedTernary._mul(self.digits, b.digits))\n",
    " \n",
    " \n",
    "def main():\n",
    "    a = BalancedTernary(\"+-0++0+\")\n",
    "    print \"a:\", a.to_int(), a\n",
    " \n",
    "    b = BalancedTernary(10)\n",
    "    print \"b:\", b.to_int(), b\n",
    " \n",
    "    c = BalancedTernary(\"+-++-\")\n",
    "    print \"c:\", c.to_int(), c\n",
    " \n",
    "    r = a * (b - c)\n",
    "    print \"a * (b - c):\", r.to_int(), r \n",
    "main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}