itertools.combinationsを無限長イテレータに対応させる

pythonのitertools.combinationsやitertools.permutationには無限長のイテレータを渡す事はできません。内部でイテレータをタプルに変換しようとするのでフリーズしてしまいます。

無限長にも対応する、アルゴリズム自体は難しくはありません。

そこで、python-ml-jpで、なぜitertools.combinationsは無限長に対応していないか質問したところ、

タプルの代わりにリストを利用することで、最初に各イテレータから全要素を取り出すことを
やめられるでしょうが、それだとreallocの回数分パフォーマンスが悪くなります。

実際に、無限長のイテレータに対応したバージョンを作ってみました。

>>> from itertools2 import combinations #無限長に対応したバージョン
>>> from itertools import count, islice
>>> it = combinations(count(), 2)       #無限長のを渡してもフリーズしない
>>> list(islice(it, 30))
[(0, 1), (0, 2), (1, 2), (0, 3), (1, 3), (2, 3), (0, 4), (1, 4), (2, 4), (3, 4),
 (0, 5), (1, 5), (2, 5), (3, 5), (4, 5), (0, 6), (1, 6), (2, 6), (3, 6), (4, 6),
 (5, 6), (0, 7), (1, 7), (2, 7), (3, 7), (4, 7), (5, 7), (6, 7), (0, 8), (1, 8)]

そして、時間測定

from __future__ import division, print_function, unicode_literals

import itertools2
import itertools

import timeit
for mod in ["itertools", "itertools2"]:
    print(mod)
    for n, r in [(10, 2), (100, 2), (1000, 2)]:
        t = timeit.Timer(
            "for x in combinations(xrange({}), {}):pass".format(n, r),
            "from {} import combinations".format(mod)
        )
        print(t.timeit(number=100))
    print()

測定結果

itertools
0.000719923900943
0.0634376715476
6.90153106051

itertools2
0.000939225516092
0.083500556635
9.52643394621

やっぱり無限長に対応した物の方が遅いです。1.5倍ぐらいずつかかっています。itertoolsのような基礎的なモジュールで50%の速度差は致命的です。しかし、単に私のコードが下手だからという可能性はぬぐえません。

無限シーケンスにcombinationsを使う機会は、実際どれくらいあるものなんでしょう？

以下、ソース。初めて作ったＣ拡張を作りました。参照関係のデバッグは地獄です。

/* itertools2.c */

#include "Python.h"
#include "structmember.h"

/* combinations object ************************************************************/

typedef struct
{
    PyObject_HEAD

    PyObject *it;
    PyObject  *pool;            /* input converted to a tuple */
    PyObject *elem;
    Py_ssize_t *indices;        /* one index per result element */
    Py_ssize_t r;               /* size of result tuple */
    Py_ssize_t i;               /* current axis of indices */
    Py_ssize_t first;
    int stopped;                /* set to 1 when the combinations iterator is exhausted */
}
combinationsobject;

static PyTypeObject combinations_type;

int next_indices(Py_ssize_t r, Py_ssize_t n, Py_ssize_t *indices, Py_ssize_t *i)
{
    if (r == 0 || n == 0 || *i < 0 || n <= *i || r < n)
    {
        return 0;
    }
    while(1)
    {
        indices[*i] += 1;
        if(*i == n - 1 && r <= indices[*i])
        {
            return 0;
        }
        else if (*i < n - 1 && indices[*i + 1] <= indices[*i])
        {
            *i += 1;
        }
        else if (0 < *i)
        {
            *i -= 1;
            indices[*i] = -1;
        }
        else
        {
            return 1;
        }
    }
}

static PyObject *
combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
    combinationsobject *co;
    Py_ssize_t r;
    PyObject *iterable = NULL;
    Py_ssize_t *indices = NULL;
    PyObject *pool = NULL;
    static char *kwargs[] = {"iterable", "r", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
                                     &iterable, &r))
    {
        return NULL;
    }

    if (r < 0)
    {
        PyErr_SetString(PyExc_ValueError, "r must be non-negative");
        goto error;
    }

    if(r > 0)
    {
        indices = PyMem_Malloc((r - 1) * sizeof(Py_ssize_t));
        if (indices == NULL)
        {
            PyErr_NoMemory();
            goto error;
        }
    }

    pool = PyList_New(0);
    if(pool == NULL)
    {
        goto error;
    }

    /* create combinationsobject structure */
    co = (combinationsobject *) type->tp_alloc(type, 0);
    if (co == NULL)
    {
        goto error;
    }

    co->pool = pool;
    co->indices = indices;
    co->it = PyObject_GetIter(iterable);
    co->first = 1;
    co->r = r;
    co->stopped = 0;

    return (PyObject *)co;
error:
    if (indices != NULL)
    {
        PyMem_Free(indices);
    }
    Py_XDECREF(pool);
    return NULL;
}

static void
combinations_dealloc(combinationsobject *co)
{
    PyObject_GC_UnTrack(co);
    Py_XDECREF(co->pool);
    Py_XDECREF(co->it);
    Py_XDECREF(co->elem);

    if (co->indices != NULL)
    {
        PyMem_Free(co->indices);
    }
    Py_TYPE(co)->tp_free(co);
}

static int
combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
{
    Py_VISIT(co->it);
    Py_VISIT(co->pool);
    Py_VISIT(co->elem);
    return 0;
}

static PyObject *
combinations_next(combinationsobject *co)
{
    Py_ssize_t j;
    PyObject *result;
    Py_ssize_t *indices = co->indices;
    Py_ssize_t *i = &co->i;
    PyObject *pool = co->pool;
    PyObject *it = co->it;
    Py_ssize_t r = co->r;

    if (co->stopped)
    {
        return NULL;
    }
    else if (r == 0)
    {
        co->first = 0;
        co->stopped = 1;
        result = PyTuple_New(0);
        if (result == NULL)
        {
            goto empty;
        }
        return result;
    }
    else if(r == 1)
    {
        PyObject *x = NULL;
        co->first = 0;
        x = PyIter_Next(it);
        if(x == NULL)
        {
            goto empty;
        }
        result = PyTuple_New(1);
        if (result == NULL)
        {
            goto empty;
        }
        PyTuple_SetItem(result, 0, x);
        return result;
    }
    else
    {
        if (co->first)
        {
            PyObject *x = NULL;
            co->first = 0;
            for(j=0; j < r - 1; ++j)
            {
                x = PyIter_Next(it);
                if(x == NULL)
                {
                    goto empty;
                }
                if (PyList_Append(pool, x) == -1)
                {
                    goto empty;
                }
                Py_DECREF(x);
            }
            x = PyIter_Next(it);
            if(x == NULL)
            {
                goto empty;
            }
            co->elem = x;
            for (j=0 ; j < r - 1  ; j++)
            {
                indices[j] = -1;
            }
            *i = r - 2;
        }
        while(1)
        {
            if(next_indices(PyList_Size(pool), r - 1, indices, i))
            {
                PyObject *x;
                result = PyTuple_New(r);
                if (result == NULL)
                {
                    goto empty;
                }
                for(j=0; j < r - 1; ++j)
                {
                    x = PyList_GetItem(pool, indices[j]);
                    if(x == NULL)
                    {
                        goto empty;
                    }
                    Py_INCREF(x);
                    PyTuple_SetItem(result, j, x);
                }
                Py_INCREF(co->elem);
                PyTuple_SetItem(result, r - 1, co->elem);
                return result;
            }
            else
            {
                if (PyList_Append(pool, co->elem) == -1)
                {
                    goto empty;
                }
                Py_DECREF(co->elem);
                co->elem = PyIter_Next(it);
                if(co->elem == NULL)
                {
                    goto empty;
                }
                for (j=0 ; j < r - 1  ; j++)
                {
                    indices[j] = -1;
                }
                *i = r - 2;
            }
        }
    }
empty:
    //Py_XDECREF(result);
    co->stopped = 1;
    return NULL;
}

PyDoc_STRVAR(combinations_doc,
             "combinations(iterable, r) --> combinations object\n\
             \n\
             Return successive r-length combinations of elements in the iterable.\n\n\
             combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");

static PyTypeObject combinations_type = {
    PyVarObject_HEAD_INIT(NULL, 0)
    "itertools2.combinations",                   /* tp_name */
    sizeof(combinationsobject),         /* tp_basicsize */
    0,                                  /* tp_itemsize */
    /* methods */
    (destructor)combinations_dealloc,           /* tp_dealloc */
    0,                                  /* tp_print */
    0,                                  /* tp_getattr */
    0,                                  /* tp_setattr */
    0,                                  /* tp_compare */
    0,                                  /* tp_repr */
    0,                                  /* tp_as_number */
    0,                                  /* tp_as_sequence */
    0,                                  /* tp_as_mapping */
    0,                                  /* tp_hash */
    0,                                  /* tp_call */
    0,                                  /* tp_str */
    PyObject_GenericGetAttr,            /* tp_getattro */
    0,                                  /* tp_setattro */
    0,                                  /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
    Py_TPFLAGS_BASETYPE,            /* tp_flags */
    combinations_doc,                           /* tp_doc */
    (traverseproc)combinations_traverse,        /* tp_traverse */
    0,                                  /* tp_clear */
    0,                                  /* tp_richcompare */
    0,                                  /* tp_weaklistoffset */
    PyObject_SelfIter,                  /* tp_iter */
    (iternextfunc)combinations_next,            /* tp_iternext */
    0,                                  /* tp_methods */
    0,                                  /* tp_members */
    0,                                  /* tp_getset */
    0,                                  /* tp_base */
    0,                                  /* tp_dict */
    0,                                  /* tp_descr_get */
    0,                                  /* tp_descr_set */
    0,                                  /* tp_dictoffset */
    0,                                  /* tp_init */
    0,                                  /* tp_alloc */
    combinations_new,                           /* tp_new */
    PyObject_GC_Del,                    /* tp_free */
};



/* module level code ********************************************************/

PyDoc_STRVAR(module_doc,
             "");


static PyMethodDef module_methods[] = {
  {NULL, NULL}           /* sentinel */
};

PyMODINIT_FUNC
inititertools2(void)
{
    int i;
    PyObject *m;
    char *name;
    PyTypeObject *typelist[] = {
       &combinations_type,
       NULL
   };

    m = Py_InitModule3("itertools2", module_methods, module_doc);
    if (m == NULL)
        return;

    for (i=0 ; typelist[i] != NULL ; i++)
    {
        if (PyType_Ready(typelist[i]) < 0)
            return;
        name = strchr(typelist[i]->tp_name, '.');
        assert (name != NULL);
        Py_INCREF(typelist[i]);
        PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
    }
}