cube-game-plus-plus/common/utils/KDTree.cpp

#include <algorithm>
#include <iostream>

#include "common/utils/KDTree.h"

KDTree::Triangle::Triangle(const Vector3& a, const Vector3& b, const Vector3& c) {
    v[0] = a;
    v[1] = b;
    v[2] = c;
    mid = (a + b + c) * (1.0f / 3.0f);
}

const Array<Vector3, 3>& KDTree::Triangle::data() const {
    return v;
}

const Vector3& KDTree::Triangle::operator[](int index) const {
    return v[index];
}

const Vector3& KDTree::Triangle::getMid() const {
    return mid;
}

KDTree::Node::Node() : splitDim(0), splitValue(0.0f), color(0xFFFFFF) {
    childs[0] = nullptr;
    childs[1] = nullptr;
}

KDTree::~KDTree() {
    clean(&root);
}

void KDTree::clean(Node* n) {
    if(n->childs[0] != nullptr) {
        clean(n->childs[0]);
    }
    if(n->childs[1] != nullptr) {
        clean(n->childs[1]);
    }
    delete n->childs[0];
    delete n->childs[1];
}

void KDTree::build(std::vector<KDTree::Triangle>& data) {
    if(data.size() == 0) {
        return;
    }
    min = data[0][0];
    max = data[0][0];
    for(const Triangle& t : data) {
        for(const Vector3& v : t.data()) {
            min.set(std::min(min[0], v[0]), std::min(min[1], v[1]), std::min(min[2], v[2]));
            max.set(std::max(max[0], v[0]), std::max(max[1], v[1]), std::max(max[2], v[2]));
        }
    }
    build(&root, data);
}

float KDTree::median(std::vector<KDTree::Triangle>& data, int dim) const {
    auto compare = [dim](const Triangle& a, const Triangle & b) {
        return a.getMid()[dim] < b.getMid()[dim];
    };
    size_t length = data.size();
    if((length & 1) == 0) {
        std::nth_element(data.begin(), data.begin() + (length / 2 - 1), data.end(), compare);
        float tmp = data[length / 2 - 1].getMid()[dim];
        std::nth_element(data.begin(), data.begin() + (length / 2), data.end(), compare);
        return (tmp + data[length / 2].getMid()[dim]) / 2;
    }
    std::nth_element(data.begin(), data.begin() + (length / 2), data.end(), compare);
    return data[length / 2].getMid()[dim];
}

void KDTree::build(Node* n, std::vector<KDTree::Triangle>& data) {
    for(KDTree::Triangle& t : data) {
        n->data.push_back(t);
    }
    if(data.size() == 0) {
        return;
    } else if(data.size() == 1) {
        n->data.push_back(data[0]);
        return;
    }
    // find min and max coordinates
    Vector3 min = data[0][0];
    Vector3 max = data[0][0];
    for(const Triangle& t : data) {
        for(const Vector3& v : t.data()) {
            min.set(std::min(min[0], v[0]), std::min(min[1], v[1]), std::min(min[2], v[2]));
            max.set(std::max(max[0], v[0]), std::max(max[1], v[1]), std::max(max[2], v[2]));
        }
    }
    // find biggest span and its dimension
    int splitDim = 0;
    float maxSpan = max[0] - min[0];
    for(int i = 1; i < 3; i++) {
        float span = max[i] - min[i];
        if(span > maxSpan) {
            splitDim = i;
            maxSpan = span;
        }
    }
    // assign data to node
    n->splitDim = splitDim;
    n->splitValue = median(data, splitDim);
    // storage for split data
    std::vector<KDTree::Triangle> lessEqualData;
    std::vector<KDTree::Triangle> greaterData;
    // actually split the data
    for(const Triangle& t : data) {
        // count points on each split side
        int lessEqualCounter = 0;
        int greaterCount = 0;
        for(const Vector3& v : t.data()) {
            if(v[n->splitDim] <= n->splitValue) {
                lessEqualCounter++;
            } else {
                greaterCount++;
            }
        }
        // put the data in the correct container
        if(lessEqualCounter == 3) {
            lessEqualData.push_back(t);
        } else if(greaterCount == 3) {
            greaterData.push_back(t);
        } else {
            n->data.push_back(t);
        }
    }
    if(lessEqualData.size() == 0 || greaterData.size() == 0) {
        for(KDTree::Triangle& t : lessEqualData) {
            n->data.push_back(t);
        }
        for(KDTree::Triangle& t : greaterData) {
            n->data.push_back(t);
        }
        return;
    }
    // recursive calls
    if(lessEqualData.size() > 0) {
        n->childs[0] = new Node();
        build(n->childs[0], lessEqualData);
    }
    if(greaterData.size() > 0) {
        n->childs[1] = new Node();
        build(n->childs[1], greaterData);
    }
}

void KDTree::fillLines(Lines& lines) const {
    lines.clear();
    lines.add(Vector3(min[0], min[1], min[2]), Vector3(max[0], min[1], min[2]), root.color);
    lines.add(Vector3(min[0], min[1], min[2]), Vector3(min[0], min[1], max[2]), root.color);
    lines.add(Vector3(max[0], min[1], min[2]), Vector3(max[0], min[1], max[2]), root.color);
    lines.add(Vector3(min[0], min[1], max[2]), Vector3(max[0], min[1], max[2]), root.color);

    lines.add(Vector3(min[0], min[1], min[2]), Vector3(min[0], max[1], min[2]), root.color);
    lines.add(Vector3(max[0], min[1], min[2]), Vector3(max[0], max[1], min[2]), root.color);
    lines.add(Vector3(min[0], min[1], max[2]), Vector3(min[0], max[1], max[2]), root.color);
    lines.add(Vector3(max[0], min[1], max[2]), Vector3(max[0], max[1], max[2]), root.color);

    lines.add(Vector3(min[0], max[1], min[2]), Vector3(max[0], max[1], min[2]), root.color);
    lines.add(Vector3(min[0], max[1], min[2]), Vector3(min[0], max[1], max[2]), root.color);
    lines.add(Vector3(max[0], max[1], min[2]), Vector3(max[0], max[1], max[2]), root.color);
    lines.add(Vector3(min[0], max[1], max[2]), Vector3(max[0], max[1], max[2]), root.color);

    fillLines(lines, &root, min, max);

    lines.build();
}

void KDTree::fillLines(Lines& lines, const Node* n, const Vector3& min, const Vector3& max) const {
    if(n->childs[0] == nullptr && n->childs[1] == nullptr) {
        return;
    }
    switch(n->splitDim) {
        case 0:
            lines.add(Vector3(n->splitValue, min[1], min[2]), Vector3(n->splitValue, max[1], min[2]), n->color);
            lines.add(Vector3(n->splitValue, max[1], min[2]), Vector3(n->splitValue, max[1], max[2]), n->color);
            lines.add(Vector3(n->splitValue, max[1], max[2]), Vector3(n->splitValue, min[1], max[2]), n->color);
            lines.add(Vector3(n->splitValue, min[1], max[2]), Vector3(n->splitValue, min[1], min[2]), n->color);
            if(n->childs[0] != nullptr) {
                fillLines(lines, n->childs[0], min, Vector3(n->splitValue, max[1], max[2]));
            }
            if(n->childs[1] != nullptr) {
                fillLines(lines, n->childs[1], Vector3(n->splitValue, min[1], min[2]), max);
            }
            break;
        case 1:
            lines.add(Vector3(min[0], n->splitValue, min[2]), Vector3(max[0], n->splitValue, min[2]), n->color);
            lines.add(Vector3(min[0], n->splitValue, min[2]), Vector3(min[0], n->splitValue, max[2]), n->color);
            lines.add(Vector3(max[0], n->splitValue, min[2]), Vector3(max[0], n->splitValue, max[2]), n->color);
            lines.add(Vector3(min[0], n->splitValue, max[2]), Vector3(max[0], n->splitValue, max[2]), n->color);
            if(n->childs[0] != nullptr) {
                fillLines(lines, n->childs[0], min, Vector3(max[0], n->splitValue, max[2]));
            }
            if(n->childs[1] != nullptr) {
                fillLines(lines, n->childs[1], Vector3(min[0], n->splitValue, min[2]), max);
            }
            break;
        case 2:
            lines.add(Vector3(min[0], min[1], n->splitValue), Vector3(min[0], max[1], n->splitValue), n->color);
            lines.add(Vector3(min[0], max[1], n->splitValue), Vector3(max[0], max[1], n->splitValue), n->color);
            lines.add(Vector3(max[0], max[1], n->splitValue), Vector3(max[0], min[1], n->splitValue), n->color);
            lines.add(Vector3(max[0], min[1], n->splitValue), Vector3(min[0], min[1], n->splitValue), n->color);
            if(n->childs[0] != nullptr) {
                fillLines(lines, n->childs[0], min, Vector3(max[0], max[1], n->splitValue));
            }
            if(n->childs[1] != nullptr) {
                fillLines(lines, n->childs[1], Vector3(min[0], min[1], n->splitValue), max);
            }
            break;
    }
}

void KDTree::cleanVisited(Node* n) {
    if(n->childs[0] != nullptr) {
        cleanVisited(n->childs[0]);
    }
    if(n->childs[1] != nullptr) {
        cleanVisited(n->childs[1]);
    }
    n->color = 0xFFFFFF;
}

bool KDTree::findIntersection(const Vector3& pos, const Vector3& direction) {
    minDistance = 9999999.0f;
    original = pos;
    cleanVisited(&root);

    float t = 0.0f;
    for(int i = 0; i < 3; i++) {
        if(direction[i] == 0.0f) {
            continue;
        }
        t = std::max(t, std::abs((min[i] - pos[i]) / direction[i]));
        t = std::max(t, std::abs((max[i] - pos[i]) / direction[i]));
    }

    return findIntersection(&root, pos, direction, t);
}

Vector3 KDTree::getIntersection() const {
    return intersection;
}

KDTree::Triangle KDTree::getIntersectedTriangle() const {
    return intersectionTriangle;
}

bool KDTree::findIntersection(Node* n, const Vector3& pos, const Vector3& direction, float tMax) {
    if(n == nullptr) {
        return false;
    }
    n->color = 0x00FFFF;
    int dim = n->splitDim;
    // lessEqual = 0
    // greater = 1
    bool r = false;
    bool first = pos[dim] > n->splitValue;
    if(direction[dim] == 0.0f) {
        r = findIntersection(n->childs[first], pos, direction, tMax);
    } else {
        float t = (n->splitValue - pos[dim]) / direction[dim];
        if(t >= 0.0f && t < tMax) {
            r = findIntersection(n->childs[first], pos, direction, t);
            r = findIntersection(n->childs[!first], pos + t * direction, direction, tMax - t) || r;
        } else {
            r = findIntersection(n->childs[first], pos, direction, tMax);
        }
    }
    
    for(KDTree::Triangle& tri : n->data) {
        float t = testIntersection(pos, direction, tri);
        if(t < 0.0f) {
            continue;
        }
        Vector3 i = pos + t * direction;
        float distance = static_cast<Vector3> (original - i).squareLength();
        if(distance > minDistance) {
            continue;
        }
        intersectionTriangle = tri;
        intersection = i;
        minDistance = distance;
        r = true;
    }
    return r;
}

float KDTree::orient(const Vector3& a, const Vector3& b, const Vector3& c, const Vector3& d) const {
    return static_cast<Vector3> (a - d).dot(static_cast<Vector3> (b - d).cross(c - d));
}

float KDTree::testIntersection(const Vector3& pos, const Vector3& direction, const KDTree::Triangle& t) const {
    const float eps = 0.0001f;
    Vector3 h1 = t[1] - t[0];
    Vector3 h2 = t[2] - t[0];
    Vector3 abc = h1.cross(h2);
    if(abc.squareLength() < eps) {
        return -1.0f;
    }
    abc.normalize();
    float d = -abc.dot(t[0]);

    float check = abc.dot(direction);
    if(std::abs(check) < eps) {
        return -1.0f;
    }
    float factor = -(abc.dot(pos) + d) / check;
    Vector3 inter = pos + factor * direction;

    int sideA = 0;
    int sideB = 0;

    for(int i = 0; i < 3; i++) {
        const Vector3& a = t[i];
        const Vector3& b = t[(i + 1) % 3];
        const Vector3 c = t[i] + abc;
        float o = orient(a, b, c, inter);
        if(o < -eps) {
            sideA++;
        } else if(o > eps) {
            sideB++;
        } else {
            sideA++;
            sideB++;
        }
    }

    if(sideA >= 3 || sideB >= 3) {
        return factor;
    }
    return -1.0f;
}