Map<String, List<JavaClass>> testClassesBySimpleClassName = new HashMap<>();
Map<String, JavaClass> classesByPackageName = new HashMap<>();

@Override
public void init(Collection<JavaClass> allClasses) {
    testClassesBySimpleClassName = allClasses.stream()
        .filter(clazz -> clazz.getName().endsWith(testClassSuffix))
        .collect(groupingBy(JavaClass::getSimpleName));

    classesByPackageName = allClasses.stream()
         .collect(toMap(JavaClass::getFullName, Function.identity()));
}

possibleTestClasses.stream()
    .map(testClazz -> classesByPackageName.getOrDefault(createImplClassNameFromTestClassName(testClazz), null))
    .noneMatch(Objects::nonNull);

private String createImplClassNameFromTestClassName(JavaClass testClazz) {
    return testClazz.getPackageName() + "." + testClazz.getSimpleName().replace(testClassSuffix, "");
}

private static ArchCondition<JavaClass> resideInTheSamePackageAsTheirTestClasses(String testClassSuffix) {
        return new ArchCondition<JavaClass>("reside in the same package as their test classes") {
            Map<String, List<JavaClass>> testClassesBySimpleClassName = new HashMap<>();
            Map<String, JavaClass> classesByPackageName = new HashMap<>();

            @Override
            public void init(Collection<JavaClass> allClasses) {
                testClassesBySimpleClassName = allClasses.stream()
                        .filter(clazz -> clazz.getName().endsWith(testClassSuffix))
                        .collect(groupingBy(JavaClass::getSimpleName));

                classesByPackageName = allClasses.stream()
                        .collect(toMap(JavaClass::getFullName, Function.identity()));
            }

            @Override
            public void check(JavaClass implementationClass,
                              ConditionEvents events) {
                String implementationClassName = implementationClass.getSimpleName();
                String implementationClassPackageName = implementationClass.getPackageName();
                String possibleTestClassName = implementationClassName + testClassSuffix;
                List<JavaClass> possibleTestClasses = testClassesBySimpleClassName.getOrDefault(possibleTestClassName,
                        emptyList());

                boolean isTestClassInWrongPackage = !possibleTestClasses.isEmpty()
                        && possibleTestClasses.stream().noneMatch(
                        clazz -> clazz.getPackageName().equals(implementationClassPackageName))
                        && implementationNotFoundInTestClassPackages(possibleTestClasses);

                if (isTestClassInWrongPackage) {
                    possibleTestClasses.forEach(wrongTestClass -> {
                        String message = createMessage(wrongTestClass,
                                String.format("does not reside in same package as implementation class <%s>",
                                        implementationClass.getName()));
                        events.add(violated(wrongTestClass, message));
                    });
                }
            }

            private boolean implementationNotFoundInTestClassPackages(List<JavaClass> possibleTestClasses) {
                return possibleTestClasses.stream()
                        .map(testClazz -> classesByPackageName.getOrDefault(
                                resolveImplClassNameFromTestClassName(testClazz), null)
                        )
                        .noneMatch(Objects::nonNull);
            }

            private String resolveImplClassNameFromTestClassName(JavaClass testClazz) {
                return testClazz.getPackageName() + "." + testClazz.getSimpleName().replace(testClassSuffix, "");
            }

        };
    }