銀月の符号

Python 使い見習いの日記・雑記

イテレータブルのある要素より後ろの要素を取り出す2

Python「インデクサとは違うのだよ、インデクサとは!」 - つまみ食う リベンジ。どうやら意図するところが「ある要素が複数あった場合、最後に現れたものの後ろ」のようで誤解していました。「最初に現れたものの後ろ」ではないのですね。

これは、どうしたものか…。

オリジナルの改良

最初の存在チェック L.count(find) は途中の探索 L.index(find) と処理内容が被っているので取り除きました。代わりに L.index が要素を見つからなかった時に送出する ValueError 例外を捕まえます。

def later0(find, iterable):
    L = list(iterable)
    L.reverse()
    try:
        L = L[:L.index(find)]
    except ValueError:
        return []
    L.reverse()
    return L

ちなみに、もし存在チェックを行うのであれば L.count(find) より find in L のほうがベターです。

別案1、ある要素で区切って、一番最後の塊が答え

takewhile を繰り返し使ってある要素で区切ります。そうしてできたリストのうち、一番最後のものが答えとなります。

個人的には気に入っているのですが、速度には難があります。リストのスライスは C で書かれている*1ので Python レベルでシーケンスをひとつずつ分割する処理を書くよりも速いのです。

from itertools import takewhile

def isplit(find, iterable):
    it = iter(iterable)
    predicate = lambda x: x != find
    while True:
        L = list(takewhile(predicate, it))
        if L:
            yield L
        else:
            break

def later1(find, iterable):
    for result in isplit(find, iterable):
        pass

    return result

別案2、先頭からある要素のインデックスを調べてスライス

enumerate はイテレータから取り出される値にインデックスを付与して (0, item0), (1, item1), (2, item2) のようにしてくれます。これを利用して、一番後ろのある要素のインデックスを求めます。あとはスライスで後部を分割します。

def later2(find, iterable):
    L = iterable if isinstance(iterable, list) else list(iterable)
    try:
        i = max(num for num, item in enumerate(L) if item == find)
    except ValueError:
        return []
    return L[i+1:]

しかし、後尾から探したほうが、見つかった際に探索を打ち切ることができるので効率がよいです。

別案3、後尾からある要素のインデックスを調べてスライス

len と xrange で逆順のインデックスを作ります。これを用いて後ろからアクセスしある要素を探索します。ある要素が見つかったら、その時のインデックスを用いてスライスし、後部を分割します。

入力によって、どのコードが速いかは変わるのですが、オリジナルの later や later0 よりも、この later3 のほうが速いことが多いです。

def later3(find, iterable):
    L = iterable if isinstance(iterable, list) else list(iterable)
    for i in xrange(len(L) - 1, -1, -1):
        if L[i] == find:
            return L[i+1:]
    else:
        return []

おまけ

countdown です。

countdown.py

普通に使うのであれば最初の4行だけで十分です。その下は C 言語によるコード用です。

def countdown(n=0):
    while True:
        yield n
        n -= 1

try:
    import _countdown
    countdown = _countdown.countdown
except ImportError:
    pass
setup.py
from distutils.core import setup, Extension

py_modules = ['countdown']
ext_modules = [
        Extension('_countdown', sources=['countdown.c'])]

setup(
    name='countdown',
    py_modules=py_modules,
    ext_modules=ext_modules)
countdown.c

itertools.count に少し手を入れただけのものです。 Python だと 4 行なのに Python/C API でまじめに書くとすごい長さですね。

#include "Python.h"

typedef struct {
    PyObject_HEAD
    Py_ssize_t cnt;
    PyObject *long_cnt;    /* Arbitrarily large count when cnt >= PY_SSIZE_T_MIN */
} countobject;

static PyTypeObject count_type;

static PyObject *
count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
    countobject *lz;
    Py_ssize_t cnt = 0;
    PyObject *cnt_arg = NULL;
    PyObject *long_cnt = NULL;

    if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
        return NULL;

    if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
        return NULL;

    if (cnt_arg != NULL) {
        cnt = PyInt_AsSsize_t(cnt_arg);
        if (cnt == -1 && PyErr_Occurred()) {
            PyErr_Clear();
            if (!PyLong_Check(cnt_arg)) {
                PyErr_SetString(PyExc_TypeError, "an integer is required");
                return NULL;
            }
            long_cnt = cnt_arg;
            Py_INCREF(long_cnt);
            cnt = PY_SSIZE_T_MIN;
        }
    }

    /* create countobject structure */
    lz = (countobject *)PyObject_New(countobject, &count_type);
    if (lz == NULL) {
        Py_XDECREF(long_cnt);
        return NULL;
    }
    lz->cnt = cnt;
    lz->long_cnt = long_cnt;

    return (PyObject *)lz;
}

static void
count_dealloc(countobject *lz)
{
    Py_XDECREF(lz->long_cnt); 
    PyObject_Del(lz);
}

static PyObject *
count_nextlong(countobject *lz)
{
    static PyObject *one = NULL;
    PyObject *cnt;
    PyObject *stepped_up;

    if (lz->long_cnt == NULL) {
        lz->long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MIN);
        if (lz->long_cnt == NULL)
            return NULL;
    }
    if (one == NULL) {
        one = PyInt_FromLong(1);
        if (one == NULL)
            return NULL;
    }
    cnt = lz->long_cnt;
    assert(cnt != NULL);
    stepped_up = PyNumber_Subtract(cnt, one);
    if (stepped_up == NULL)
        return NULL;
    lz->long_cnt = stepped_up;
    return cnt;
}

static PyObject *
count_next(countobject *lz)
{
    if (lz->cnt == PY_SSIZE_T_MIN)
        return count_nextlong(lz);
    return PyInt_FromSsize_t(lz->cnt--);
}

static PyObject *
count_repr(countobject *lz)
{
    PyObject *cnt_repr;
    PyObject *result;

    if (lz->cnt != PY_SSIZE_T_MIN)
        return PyString_FromFormat("count(%zd)", lz->cnt);

    cnt_repr = PyObject_Repr(lz->long_cnt);
    if (cnt_repr == NULL)
        return NULL;
    result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
    Py_DECREF(cnt_repr);
    return result;
}

PyDoc_STRVAR(count_doc,
"countdown([firstval]) --> countdown object\n\
\n\
Return a countdown object whose .next() method returns consecutive\n\
integers starting from zero or, if specified, from firstval.");

static PyTypeObject count_type = {
    PyVarObject_HEAD_INIT(NULL, 0)
    "_countdown.countdown",  /* tp_name */
    sizeof(countobject),  /* tp_basicsize */
    0,                /* tp_itemsize */
    /* methods */
    (destructor)count_dealloc, /* tp_dealloc */
    0,                /* tp_print */
    0,                /* tp_getattr */
    0,                /* tp_setattr */
    0,                /* tp_compare */
    (reprfunc)count_repr, /* 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, /* tp_flags */
    count_doc,        /* tp_doc */
    0,                /* tp_traverse */
    0,                /* tp_clear */
    0,                /* tp_richcompare */
    0,                /* tp_weaklistoffset */
    PyObject_SelfIter, /* tp_iter */
    (iternextfunc)count_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 */
    count_new,        /* tp_new */

};

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

PyMODINIT_FUNC
init_countdown(void)
{
    PyObject *m;

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

    if (PyType_Ready(&count_type) < 0)
        return;
    Py_INCREF(&count_type);
    PyModule_AddObject(m, "countdown", (PyObject *)&count_type);
}

*1:CPython の場合