リストをN個に分割 - Groovy, Java, Python

Java Groovy Python

リストをなるべく均等に N 分割する処理を Groovy・Java・Python で実装してみました。

今回は、[0, 1, 2, 3, 4, 5, 6, 7] を 3分割した結果が [[0, 1, 2], [3, 4, 5], [6, 7]] となるような処理を想定しています。 （余り分を先頭から順に1つずつ分配）

ソースは http://github.com/fits/try_samples/tree/master/blog/20151109/

Groovy の場合

  • Groovy 2.4.5

畳み込みを使って実装してみました。

  • (1) サブリストとして取り出す範囲を算出
  • (2) 取り出したサブリストを結果へ追加
divide_list.groovy
def divideList = { xs, n ->
    int q = xs.size() / n
    int m = xs.size() % n

    (0..<n).inject([]) { acc, i ->
        // (1) サブリスト化する範囲の開始・終了位置を算出
        def fr = acc*.size().sum(0)
        def to = fr + q + ((i < m)? 1: 0)

        // (2) サブリストを取り出し、結果へ追加
        acc << xs[fr..<to]
    }
}

println divideList(0..<8, 3)
println divideList(0..<7, 3)
println divideList(0..<6, 3)
println divideList(0..<6, 6)
実行結果
> groovy divide_list.groovy

[[0, 1, 2], [3, 4, 5], [6, 7]]
[[0, 1, 2], [3, 4], [5, 6]]
[[0, 1], [2, 3], [4, 5]]
[[0], [1], [2], [3], [4], [5]]

なお、上記の実装内容では divideList(0..<3, 5) のように要素数よりも分割数を多くすると [[0], [1], [2], [], []] のように不足分が空リストとなります。

場合によっては、以下のように小さい数に合わせた方が実用的かもしれません。

def divideList = { xs, n ->
    ・・・
    nn = Math.min(n, xs.size())

    (0..<nn).inject([]) { acc, i ->
        ・・・
    }
}

Java 8 の場合

Groovy 版と同様に実装しました。 Stream の collect メソッドを使っています。

DivideList.java
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class DivideList {
    public static void main(String... args) {
        System.out.println(divideList(range(0, 8), 3));
        System.out.println(divideList(range(0, 7), 3));
        System.out.println(divideList(range(0, 6), 3));
        System.out.println(divideList(range(0, 6), 6));
    }

    private static <T> List<List<T>> divideList(List<T> xs, int n) {
        int q = xs.size() / n;
        int m = xs.size() % n;

        return IntStream.range(0, n).collect(
                ArrayList::new,
                (acc, i) -> {
                    int fr = acc.stream().mapToInt(List::size).sum();
                    int to = fr + q + ((i < m)? 1: 0);

                    acc.add(xs.subList(fr, to));
                },
                ArrayList::addAll
        );
    }

    private static List<Integer> range(int start, int end) {
        return IntStream.range(start, end).boxed().collect(Collectors.toList());
    }
}
実行結果
> java DivideList

[[0, 1, 2], [3, 4, 5], [6, 7]]
[[0, 1, 2], [3, 4], [5, 6]]
[[0, 1], [2, 3], [4, 5]]
[[0], [1], [2], [3], [4], [5]]

Python の場合

Python も同様に実装しました。

変数 fr に値を保持させるため、ラムダを入れ子にして引数のデフォルト値を利用しています。

また、以下のように / で割ると結果が実数になって都合が悪いため、// を使っています。

式 結果
5 / 2 2.5
5 // 2 2
divide_list.py
from functools import reduce

def divide_list(xs, n):
    q = len(xs) // n
    m = len(xs) % n

    return reduce(
        lambda acc, i:
            (lambda fr = sum([ len(x) for x in acc ]):
                acc + [ xs[fr:(fr + q + (1 if i < m else 0))] ]
            )()
        ,
        range(n),
        []
    )

range_list = lambda n: list(range(n))

print(divide_list(range_list(8), 3))
print(divide_list(range_list(7), 3))
print(divide_list(range_list(6), 3))
print(divide_list(range_list(6), 6))
実行結果
> python divide_list.py

[[0, 1, 2], [3, 4, 5], [6, 7]]
[[0, 1, 2], [3, 4], [5, 6]]
[[0, 1], [2, 3], [4, 5]]
[[0], [1], [2], [3], [4], [5]]