SmallDemo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Для чтения/записи csv файлов\n",
    "import csv\n",
    "# Для работы с архивами\n",
    "import gzip\n",
    "# Для работы с файловой системой\n",
    "import os\n",
    "# Эффективные массивы простых типов\n",
    "import numpy\n",
    "# Работа с матрицами (подсчет общих друзей реализован как умножение матрицы графа самое на себя)\n",
    "import scipy\n",
    "from scipy.sparse import coo_matrix, csr_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Пути к данным\n",
    "dataPath = \"./\"\n",
    "graphPath = os.path.join(dataPath, \"trainGraph\")\n",
    "predictionPath = os.path.join(dataPath,\"prediction.gz\")\n",
    "\n",
    "# Основные параметры графа\n",
    "numUsers = 107474\n",
    "numLinks = 72384968 / 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# В этих массивах мы будем собирать данные. Инициализируем их заранее нужным размером чтобы \n",
    "# небыло лишнего копирования\n",
    "formUser = numpy.zeros( (numLinks), dtype=numpy.int32 ) \n",
    "toUser = numpy.zeros( (numLinks), dtype=numpy.int32 ) \n",
    "data = numpy.ones( (numLinks), dtype=numpy.int16 ) \n",
    "\n",
    "# Здесь храним позицию, на которую надо записать новую связь\n",
    "current = 0\n",
    "\n",
    "# Запоминаем самы старший ид юзера\n",
    "maxUserId = 0\n",
    "\n",
    "# Итерируемся по файлам в папке\n",
    "for file in [f for f in os.listdir(graphPath) if f.startswith(\"part\")]:\n",
    "    csvinput = gzip.open(os.path.join(graphPath, file))\n",
    "    csv_reader = csv.reader(csvinput, delimiter='\\t')\n",
    "    # А теперь по строкам в файле\n",
    "    for line in csv_reader:\n",
    "        user = int(line[0]) \n",
    "        maxUserId = max(user,maxUserId)\n",
    "        # Разбираем идшки и маски друзей\n",
    "        for friendship in line[1].replace(\"{(\", \"\").replace(\")}\", \"\").split(\"),(\"):\n",
    "            parts=friendship.split(\",\")\n",
    "            # Записываем связь в массивы и двигаем указатель\n",
    "            formUser[current] = user\n",
    "            friend = int(parts[0])\n",
    "            toUser[current] = friend\n",
    "            maxUserId = max(friend,maxUserId)\n",
    "            current += 1\n",
    "            \n",
    "    # Не забываем закрыть файл\n",
    "    csvinput.close()    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Создаем из массивов матрицу. Изначальна матрица хранится в виде списка [i,j,v], но для эффективной \n",
    "# дальнейшей работы нам надо преобразовать в вид [i->[j,v]]\n",
    "fullMatrix = coo_matrix(\n",
    "    (data, (formUser, toUser)),\n",
    "    shape=(maxUserId + 1, maxUserId + 1)).tocsr()\n",
    "\n",
    "# Массивы больше не нужны, удаляем их из памяти\n",
    "del formUser\n",
    "del toUser\n",
    "del data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Считаем транспонированную матрицу (колонки и ряды поменяны местами) и её тоже приводим в вид [i->[j,v]]\n",
    "reversedMatrix = scipy.transpose(fullMatrix).tocsr()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Поскольку прогноз нам нужно строить только для части пользователей, остальных из\n",
    "# исходной матрицы уберем (забьем нулями)\n",
    "for i in range(maxUserId + 1):\n",
    "    if i % 11 != 7:\n",
    "        ptr = fullMatrix.indptr[i]\n",
    "        ptr_next = fullMatrix.indptr[i+1]\n",
    "        if ptr != ptr_next:\n",
    "            fullMatrix.data[ptr:ptr_next].fill(0)\n",
    "            \n",
    "# Чтобы нули не мешались при умножении, вычистим их и подуменьшим матрицу\n",
    "fullMatrix.eliminate_zeros()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Здесь и происходит основная магия - через умножение матриц получаем счетчики общих друзей,\n",
    "# По которым сделаем прогноз\n",
    "commonFriends = fullMatrix.dot(reversedMatrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Теперь осталось его записать в файл. Открываем врайтеры\n",
    "output = gzip.open(predictionPath, \"w\")\n",
    "writer = csv.writer(output, delimiter='\\t')\n",
    "\n",
    "for i in range(maxUserId + 1):\n",
    "    # Два указателя дают нам границы в которых лежат данные для этого i в матрице\n",
    "    ptr = commonFriends.indptr[i]\n",
    "    ptr_next = commonFriends.indptr[i+1] \n",
    "    # Если они не равны, значит данные есть и можно экспортировать\n",
    "    if ptr != ptr_next:\n",
    "        # Достаем счетчики общих друзей и создаем порядок на них от большего к меньшему\n",
    "        counts = commonFriends.data[ptr:ptr_next]\n",
    "        order = numpy.argsort(-counts)\n",
    "        \n",
    "        # Не забываем что из прогноза надо убрать себя и своих известных друзей\n",
    "        mineFriends = set(fullMatrix.indices[fullMatrix.indptr[i]:fullMatrix.indptr[i+1]])\n",
    "        mineFriends.add(i)\n",
    "        \n",
    "        # Достаем идшки друзей, сортируем, фильтруем, обрезаем и пишем\n",
    "        ids = commonFriends.indices[ptr:ptr_next]        \n",
    "        writer.writerow([i] + filter(lambda x: x not in mineFriends, ids[order])[:42])      \n",
    "\n",
    "# Не забываем закрыть файл\n",
    "output.close()        "
   ]
  }
 ],
 "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.5.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}