Itertools for efficient looping.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python Itertools - 高效的循环\n",
    "作者: tushushu  \n",
    "项目地址: https://github.com/tushushu/flying-python  \n",
    "\n",
    "Python官方文档用\"高效的循环\"来形容itertools模块，有些tools会带来性能提升，而另外一些tools并不快，只是会节省一些开发时间而已，如果滥用还会导致代码可读性变差。我们不妨把itertools的兄弟们拉出来溜溜。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 数列累加\n",
    "给定一个列表An，返回数列累加和Sn。\n",
    "举例说明：\n",
    "* 输入: [1, 2, 3, 4, 5]\n",
    "* 返回: [1, 3, 6, 10, 15]  \n",
    "\n",
    "使用accumulate，性能提升了2.5倍"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import accumulate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _accumulate_list(arr):\n",
    "    tot = 0\n",
    "    for x in arr:\n",
    "        tot += x\n",
    "        yield tot\n",
    "\n",
    "def accumulate_list(arr):\n",
    "    return list(_accumulate_list(arr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fast_accumulate_list(arr):\n",
    "    return list(accumulate(arr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = list(range(1000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "61 µs ± 2.91 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit accumulate_list(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "21.3 µs ± 811 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit fast_accumulate_list(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 选择数据\n",
    "给定一个列表data，一个用0/1表示的列表selectors，返回被选择的数据。\n",
    "举例说明：\n",
    "* 输入: [1, 2, 3, 4, 5], [0, 1, 0, 1, 0]\n",
    "* 返回: [2, 4]  \n",
    "\n",
    "使用compress，性能提升了2.8倍"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import compress\n",
    "from random import randint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_data(data, selectors):\n",
    "    return [x for x, y in zip(data, selectors) if y]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fast_select_data(data, selectors):\n",
    "    return list(compress(data, selectors))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = list(range(10000))\n",
    "selectors = [randint(0, 1) for _ in range(10000)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "341 µs ± 17.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit select_data(data, selectors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "130 µs ± 3.19 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit fast_select_data(data, selectors)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 组合\n",
    "给定一个列表arr和一个数字k，返回从arr中选择k个元素的所有情况。\n",
    "举例说明：\n",
    "* 输入: [1, 2, 3], 2\n",
    "* 返回: [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)] \n",
    "\n",
    "使用permutations，性能提升了10倍"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import permutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _get_permutations(arr, k, i):\n",
    "    if i == k:\n",
    "        return [arr[:k]]\n",
    "    res = []\n",
    "    for j in range(i, len(arr)):\n",
    "        arr_cpy = arr.copy()\n",
    "        arr_cpy[i], arr_cpy[j] = arr_cpy[j], arr_cpy[i]\n",
    "        res += _get_permutations(arr_cpy, k, i + 1)\n",
    "    return res\n",
    "    \n",
    "def get_permutations(arr, k):\n",
    "    return _get_permutations(arr, k, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fast_get_permutations(arr, k):\n",
    "    return list(permutations(arr, k))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = list(range(10))\n",
    "k = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "15.5 ms ± 1.96 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit -n 1 get_permutations(arr, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.56 ms ± 284 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit -n 1 fast_get_permutations(arr, k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 筛选数据\n",
    "给定一个列表arr，筛选出所有的偶数。\n",
    "举例说明：\n",
    "* 输入: [3, 1, 4, 5, 9, 2]\n",
    "* 返回: [(4, 2] \n",
    "\n",
    "使用filterfalse，性能反而会变慢，所以不要迷信itertools。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import filterfalse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_even_nums(arr):\n",
    "    return [x for x in arr if x % 2 == 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fast_get_even_nums(arr):\n",
    "    return list(filterfalse(lambda x: x % 2, arr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = list(range(10000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "417 µs ± 18.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit get_even_nums(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "823 µs ± 22.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit fast_get_even_nums(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. 条件终止\n",
    "给定一个列表arr，依次对列表的所有数字进行求和，若遇到某个元素大于target之后则终止求和，返回这个和。\n",
    "举例说明：\n",
    "* 输入: [1, 2, 3, 4, 5], 3\n",
    "* 返回: 6 (4 > 3，终止)\n",
    "\n",
    "使用takewhile，性能反而会变慢，所以不要迷信itertools。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import takewhile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cond_sum(arr, target):\n",
    "    res = 0\n",
    "    for x in arr:\n",
    "        if x > target:\n",
    "            break\n",
    "        res += x\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fast_cond_sum(arr, target):\n",
    "    return sum(takewhile(lambda x: x <= target, arr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = list(range(10000))\n",
    "target = 5000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "245 µs ± 11.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit cond_sum(arr, target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "404 µs ± 13.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit fast_cond_sum(arr, target)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. 循环嵌套\n",
    "给定列表arr1，arr2，返回两个列表的所有元素两两相加的和。\n",
    "举例说明：\n",
    "* 输入: [1, 2], [4, 5]\n",
    "* 返回: [1 + 4， 1 + 5， 2 + 4， 2 + 5]\n",
    "\n",
    "使用product，性能提升了1.25倍。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import product"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _cross_sum(arr1, arr2):\n",
    "    for x in arr1:\n",
    "        for y in arr2:\n",
    "            yield x + y\n",
    "\n",
    "def cross_sum(arr1, arr2):\n",
    "    return list(_cross_sum(arr1, arr2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fast_cross_sum(arr1, arr2):\n",
    "    return [x + y for x, y in product(arr1, arr2)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr1 = list(range(100))\n",
    "arr2 = list(range(100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "484 µs ± 16.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit cross_sum(arr1, arr2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "373 µs ± 11.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit fast_cross_sum(arr1, arr2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. 二维列表转一维列表\n",
    "给定二维列表arr，转为一维列表\n",
    "举例说明：\n",
    "* 输入: [[1, 2], [3, 4]]\n",
    "* 返回: [1, 2, 3, 4]\n",
    "\n",
    "使用chain，性能提升了6倍。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _flatten(arr2d):\n",
    "    for arr in arr2d:\n",
    "        for x in arr:\n",
    "            yield x\n",
    "\n",
    "def flatten(arr2d):\n",
    "    return list(_flatten(arr2d))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fast_flatten(arr2d):\n",
    "    return list(chain(*arr2d))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr2d = [[x + y * 100 for x in range(100)] for y in range(100)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "379 µs ± 15.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit flatten(arr2d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "66.9 µs ± 3.43 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit fast_flatten(arr2d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}