introduce: partition by count and by arbitrary condition

This commit is contained in:
Gregory.Shrago
2016-07-30 01:27:38 +03:00
parent 216b92b438
commit 732f6943ed
3 changed files with 175 additions and 46 deletions
@@ -29,6 +29,7 @@ import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import static com.intellij.openapi.util.Conditions.not;
import static com.intellij.util.containers.JBIterable.SeparatorOption.*;
/**
* @author gregsh
@@ -299,6 +300,29 @@ public class TreeTraverserTest extends TestCase {
assertEquals(Arrays.asList(0, 1, 0, 2, 0, 3, 0, 4, 0, 5), it.flatMap(o -> ContainerUtil.list(0, o)).toList());
}
public void testPartition() {
JBIterable<Integer> it = JBIterable.of(1, 2, 3, 4, 5);
assertEquals(Arrays.asList(Arrays.asList(1, 2), Arrays.asList(3, 4)), it.partition(2, true).toList());
assertEquals(Arrays.asList(Arrays.asList(1, 2), Arrays.asList(3, 4), Arrays.asList(5)), it.partition(2, false).toList());
assertEquals("[[1, 2], [4, 5]]", it.partition(SKIP, o -> o % 3 == 0).map(o -> o.toList()).toList().toString());
assertEquals("[[1, 2], [3], [4, 5]]", it.partition(EXTRACT, o -> o % 3 == 0).map(o -> o.toList()).toList().toString());
assertEquals("[[1, 2, 3], [4, 5]]", it.partition(HEAD, o -> o % 3 == 0).map(o -> o.toList()).toList().toString());
assertEquals("[[1, 2], [3, 4, 5]]", it.partition(TAIL, o -> o % 3 == 0).map(o -> o.toList()).toList().toString());
assertEquals("[[1, 2, 3, 4], [5]]", it.partition(EXTRACT, o -> o == 5).map(o -> o.toList()).toList().toString());
assertEquals("[[], [1], [2, 3, 4, 5]]", it.partition(EXTRACT, o -> o == 1).map(o -> o.toList()).toList().toString());
assertEquals("[[], [], [], [], []]", it.partition(SKIP, o -> true).map(o -> o.toList()).toList().toString());
assertEquals("[[1], [2], [3], [4], [5]]", it.partition(HEAD, o -> true).map(o -> o.toList()).toList().toString());
assertEquals("[[], [1], [2], [3], [4], [5]]", it.partition(TAIL, o -> true).map(o -> o.toList()).toList().toString());
assertEquals("[[], [1], [], [2], [], [3], [], [4], [], [5]]", it.partition(EXTRACT, o -> true).map(o -> o.toList()).toList().toString());
assertEquals(3, it.partition(EXTRACT, o -> o % 3 == 0).size());
assertEquals(10, it.partition(EXTRACT, o -> true).size());
assertEquals(it.partition(2, false).toList(), it.partition(HEAD, o -> o % 2 == 0).map(o -> o.toList()).toList());
}
// TreeTraversal ----------------------------------------------
@NotNull
@@ -482,6 +482,85 @@ public abstract class JBIterable<E> implements Iterable<E> {
return map(function).flatten(Function.ID);
}
/**
* Returns a {@code JBIterable} that groups this iterable into lists of the specified size.
* If 'strict' flag is true only groups of size 'n' are returned.
*/
@NotNull
public final JBIterable<List<E>> partition(final int n, final boolean strict) {
if (n <= 0) throw new IllegalArgumentException(n + " <= 0");
return intercept(new Function<Iterator<E>, Iterator<List<E>>>() {
@Override
public Iterator<List<E>> fun(Iterator<E> iterator) {
final Iterator<E> orig = iterator;
return new JBIterator<List<E>>() {
@Override
protected List<E> nextImpl() {
ArrayList<E> next = ContainerUtil.newArrayListWithCapacity(n);
for (E e : once(orig).take(n)) {
next.add(e);
}
return next.isEmpty() || strict && next.size() < n ? stop() : next;
}
};
}
});
}
public enum SeparatorOption {HEAD, TAIL, EXTRACT, SKIP}
/**
* Returns a {@code JBIterable} that groups this iterable by the specified condition
* without additional memory allocation.
*/
@NotNull
public final JBIterable<JBIterable<E>> partition(final SeparatorOption option, final Condition<E> condition) {
return intercept(new Function<Iterator<E>, Iterator<JBIterable<E>>>() {
@Override
public Iterator<JBIterable<E>> fun(Iterator<E> iterator) {
final Iterator<E> orig = iterator;
return new JBIterator<JBIterable<E>>() {
List<E> stored;
JBIterator<E> it;
@Override
protected JBIterable<E> nextImpl() {
// iterate through the previous result fully before proceeding
if (it != null && it.hasNext()) once(it).size();
it = null;
List<E> sep = stored;
stored = null;
if (option == SeparatorOption.EXTRACT && sep != null) return JBIterable.from(sep);
if (!orig.hasNext()) {
return option == SeparatorOption.TAIL && sep != null ? JBIterable.from(sep) : stop();
}
it = JBIterator.wrap(orig);
JBIterable<E> next = once(it.takeWhile(new Condition<E>() {
@Override
public boolean value(E e) {
if (!condition.value(e)) return true;
stored = Collections.singletonList(e);
return false;
}
}));
switch (option) {
case HEAD: return next.append(new JBIterable<E>() {
@Override
public Iterator<E> iterator() {
return stored != null ? stored.iterator() : JBIterable.<E>empty().iterator();
}
});
case TAIL: return sep != null ? JBIterable.from(sep).append(next) : next;
case EXTRACT:
case SKIP: return next;
default: throw new AssertionError(option);
}
}
};
}
});
}
/**
* Determines whether this iterable is empty.
*/
@@ -60,7 +60,12 @@ public abstract class JBIterator<E> implements Iterator<E> {
@NotNull
public static <E> JBIterator<E> from(@NotNull final Iterator<E> it) {
return it instanceof JBIterator ? (JBIterator<E>)it : new JBIterator<E>() {
return it instanceof JBIterator ? (JBIterator<E>)it : wrap(it);
}
@NotNull
static <E> JBIterator<E> wrap(@NotNull final Iterator<E> it) {
return new JBIterator<E>() {
@Override
protected E nextImpl() {
return it.hasNext() ? it.next() : stop();
@@ -71,7 +76,7 @@ public abstract class JBIterator<E> implements Iterator<E> {
private Object myCurrent = NONE;
private Object myNext = NONE;
private final Op myFirstOp = new Op(null);
private Op myFirstOp = new Op(null);
private Op myLastOp = myFirstOp;
/**
@@ -141,7 +146,7 @@ public abstract class JBIterator<E> implements Iterator<E> {
if (myNext != NONE) return;
Object o = NONE;
for (Op op = myFirstOp; op != null; op = op == null ? myFirstOp : op.nextOp) {
o = op == myFirstOp ? nextImpl() : op.apply(o);
o = op.impl == null ? nextImpl() : op.apply(o);
if (myNext == SKIP) {
o = myNext = NONE;
op = null;
@@ -153,42 +158,23 @@ public abstract class JBIterator<E> implements Iterator<E> {
@NotNull
public final <T> JBIterator<T> transform(@NotNull Function<? super E, T> function) {
return addOp(new Op<Function<? super E, T>>(function) {
@Override
public Object apply(Object o) {
return impl.fun((E)o);
}
});
return addOp(true, new TransformOp<E, T>(function));
}
@NotNull
public final JBIterator<E> filter(@NotNull Condition<? super E> condition) {
return addOp(new Op<Condition<? super E>>(condition) {
@Override
public Object apply(Object o) {
return impl.value((E)o) ? o : skip();
}
});
return addOp(true, new FilterOp<E>(condition));
}
@NotNull
public final JBIterator<E> take(int count) {
return takeWhile(new CountDown<E>(count));
// add first so that the underlying iterator stay on 'count' position
return addOp(myLastOp.impl != null, new WhileOp<E>(new CountDown<E>(count)));
}
@NotNull
public final JBIterator<E> takeWhile(@NotNull Condition<? super E> condition) {
return addOp(new Op<Condition<? super E>>(condition) {
@Override
public Object apply(Object o) {
return impl.value((E)o) ? o : stop();
}
@Override
public String toString() {
return "takeWhile:" + super.toString();
}
});
return addOp(true, new WhileOp<E>(condition));
}
@NotNull
@@ -198,28 +184,19 @@ public abstract class JBIterator<E> implements Iterator<E> {
@NotNull
public final JBIterator<E> skipWhile(@NotNull final Condition<? super E> condition) {
return addOp(new Op<Condition<? super E>>(condition) {
boolean active = true;
@Override
public Object apply(Object o) {
if (active && condition.value((E)o)) return skip();
active = false;
return o;
}
@Override
public String toString() {
return "skipWhile:" + super.toString();
}
});
return addOp(true, new SkipOp<E>(condition));
}
@NotNull
private <T> T addOp(@NotNull Op op) {
myLastOp.nextOp = op;
myLastOp = myLastOp.nextOp;
private <T> T addOp(boolean last, @NotNull Op op) {
if (last) {
myLastOp.nextOp = op;
myLastOp = myLastOp.nextOp;
}
else {
op.nextOp = myFirstOp;
myFirstOp = op;
}
return (T)this;
}
@@ -301,4 +278,53 @@ public abstract class JBIterator<E> implements Iterator<E> {
return cur > 0 && cur-- != 0;
}
}
private static class TransformOp<E, T> extends Op<Function<? super E, T>> {
TransformOp(Function<? super E, T> function) {
super(function);
}
@Override
public Object apply(Object o) {
return impl.fun((E)o);
}
}
private class FilterOp<E> extends Op<Condition<? super E>> {
FilterOp(Condition<? super E> condition) {
super(condition);
}
@Override
public Object apply(Object o) {
return impl.value((E)o) ? o : skip();
}
}
private class WhileOp<E> extends Op<Condition<? super E>> {
WhileOp(Condition<? super E> condition) {
super(condition);
}
@Override
public Object apply(Object o) {
return impl.value((E)o) ? o : stop();
}
}
private class SkipOp<E> extends Op<Condition<? super E>> {
boolean active;
SkipOp(Condition<? super E> condition) {
super(condition);
active = true;
}
@Override
public Object apply(Object o) {
if (active && impl.value((E)o)) return skip();
active = false;
return o;
}
}
}