feat: implement lazy graph expansion - Task 1: hub station selection, BuildGraphFromHubs, and ExpandGraphLazy
This commit is contained in:
@@ -1,6 +1,9 @@
|
||||
package routing
|
||||
|
||||
import "sort"
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
)
|
||||
|
||||
// Edge represents a graph edge connecting two nodes.
|
||||
type Edge struct {
|
||||
@@ -43,6 +46,68 @@ const (
|
||||
EdgeKindSynthetic
|
||||
)
|
||||
|
||||
// hubCriteria defines the criteria for selecting hub stations.
|
||||
type hubCriteria struct {
|
||||
minPopulation int // minimum city population in millions to be considered a hub
|
||||
minOutgoingFlights int // minimum number of outgoing Yandex flights to be considered a hub
|
||||
defaultOutgoingFlights int // default outgoing flights count when data is unavailable
|
||||
}
|
||||
|
||||
// HubStation represents a selected hub station with its selection rationale.
|
||||
type HubStation struct {
|
||||
// Station is the underlying station node.
|
||||
Station *Node
|
||||
// CityCode is the city the station belongs to.
|
||||
CityCode string
|
||||
// OutgoingFlights is the estimated number of outgoing Yandex flights from this station.
|
||||
OutgoingFlights int
|
||||
// Population is the city population in millions used for hub selection.
|
||||
Population int
|
||||
// IsHub indicates whether this station meets the hub criteria.
|
||||
IsHub bool
|
||||
}
|
||||
|
||||
// HubStationSelectionResult holds the results of hub station selection.
|
||||
type HubStationSelectionResult struct {
|
||||
// Hubs are the selected hub stations sorted by priority.
|
||||
Hubs []*HubStation
|
||||
// Rejected are stations that don't meet hub criteria, with reason.
|
||||
Rejected []*HubStation
|
||||
}
|
||||
|
||||
// SelectHubStations selects hub stations from a list based on criteria.
|
||||
// Hubs are selected based on: population (million+ cities), number of outgoing Yandex flights.
|
||||
func SelectHubStations(stations []StationInfo, criteria hubCriteria) HubStationSelectionResult {
|
||||
result := HubStationSelectionResult{
|
||||
Hubs: []*HubStation{},
|
||||
Rejected: []*HubStation{},
|
||||
}
|
||||
|
||||
for _, si := range stations {
|
||||
hub := &HubStation{
|
||||
Station: &Node{ID: si.ID, Type: NodeTypeStation, Name: si.Name, CityCode: si.CityCode},
|
||||
CityCode: si.CityCode,
|
||||
OutgoingFlights: criteria.defaultOutgoingFlights,
|
||||
Population: 0, // will be inferred from city code later
|
||||
IsHub: false,
|
||||
}
|
||||
|
||||
// A station is considered a hub if:
|
||||
// 1. It has >= minOutgoingFlights (outgoing Yandex flight data available) - primary criterion
|
||||
// For MVP, outgoing flights is the primary criterion.
|
||||
hasOutgoingFlights := hub.OutgoingFlights >= criteria.minOutgoingFlights
|
||||
|
||||
if hasOutgoingFlights {
|
||||
hub.IsHub = true
|
||||
result.Hubs = append(result.Hubs, hub)
|
||||
} else {
|
||||
result.Rejected = append(result.Rejected, hub)
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// StationInfo holds station information for graph building from a station directory.
|
||||
type StationInfo struct {
|
||||
ID string
|
||||
@@ -146,7 +211,180 @@ func BuildGraphFromStations(stations []StationInfo) *Graph {
|
||||
return graph
|
||||
}
|
||||
|
||||
|
||||
|
||||
// BuildGraphFromHubs builds a routing graph from a list of station info records,
|
||||
// focusing on hub stations. It creates station nodes and city hub nodes with
|
||||
// synthetic edges connecting stations to their city hubs, similar to
|
||||
// BuildGraphFromStations but optimized for hub-based lazy expansion.
|
||||
func BuildGraphFromHubs(stations []StationInfo, hubCriteria hubCriteria) *Graph {
|
||||
graph := NewGraph()
|
||||
|
||||
// Select hub stations based on criteria
|
||||
selection := SelectHubStations(stations, hubCriteria)
|
||||
|
||||
// Track city nodes by code to avoid duplicates
|
||||
cityNodes := make(map[string]*Node)
|
||||
|
||||
// Add hub station nodes and create/connect city hub nodes
|
||||
for _, hub := range selection.Hubs {
|
||||
si := findStationByID(stations, hub.Station.ID)
|
||||
|
||||
// Add station node
|
||||
station := &Node{
|
||||
ID: hub.Station.ID,
|
||||
Type: NodeTypeStation,
|
||||
Name: hub.Station.Name,
|
||||
CityCode: hub.CityCode,
|
||||
}
|
||||
graph.AddNode(station)
|
||||
|
||||
// Create or retrieve city hub node
|
||||
cityKey := "city:" + hub.CityCode
|
||||
if _, exists := cityNodes[hub.CityCode]; !exists {
|
||||
cityNode := &Node{
|
||||
ID: cityKey,
|
||||
Type: NodeTypeCity,
|
||||
Name: si.CityName,
|
||||
}
|
||||
graph.AddNode(cityNode)
|
||||
cityNodes[hub.CityCode] = cityNode
|
||||
}
|
||||
|
||||
cityNode := cityNodes[hub.CityCode]
|
||||
|
||||
// Add synthetic edge: station <-> city hub
|
||||
graph.AddEdge(&Edge{
|
||||
From: station,
|
||||
To: cityNode,
|
||||
Kind: EdgeKindSynthetic,
|
||||
Duration: 300, // 5 min synthetic transfer
|
||||
Transport: "train",
|
||||
IsTransfer: true,
|
||||
})
|
||||
|
||||
// Add reverse synthetic edge: city hub -> station
|
||||
graph.AddEdge(&Edge{
|
||||
From: cityNode,
|
||||
To: station,
|
||||
Kind: EdgeKindSynthetic,
|
||||
Duration: 300, // 5 min synthetic transfer
|
||||
Transport: "train",
|
||||
IsTransfer: true,
|
||||
})
|
||||
}
|
||||
|
||||
// Also add non-hub station nodes without city connections (they'll be expanded lazily)
|
||||
for _, s := range stations {
|
||||
if !hubExists(selection.Hubs, s.ID) {
|
||||
// Add station node without city connection for lazy expansion
|
||||
station := &Node{
|
||||
ID: s.ID,
|
||||
Type: NodeTypeStation,
|
||||
Name: s.Name,
|
||||
CityCode: s.CityCode,
|
||||
}
|
||||
graph.AddNode(station)
|
||||
}
|
||||
}
|
||||
|
||||
return graph
|
||||
}
|
||||
|
||||
// SortEdges sorts edges by duration in ascending order (shortest first).
|
||||
|
||||
|
||||
// ExpandGraphLazy on-demand adds edges from the current node to hub candidates.
|
||||
// This enables graph expansion during BFS route search without pre-building the
|
||||
// complete graph, staying within API quota constraints.
|
||||
func (g *Graph) ExpandGraphLazy(currentNode *Node, destCityCode string) error {
|
||||
if currentNode == nil {
|
||||
return fmt.Errorf("currentNode cannot be nil")
|
||||
}
|
||||
switch currentNode.Type {
|
||||
case NodeTypeStation:
|
||||
return expandFromStation(g, currentNode, destCityCode)
|
||||
case NodeTypeCity:
|
||||
return expandFromCityHub(g, currentNode, destCityCode)
|
||||
default:
|
||||
return fmt.Errorf("unsupported node type: %d", currentNode.Type)
|
||||
}
|
||||
}
|
||||
|
||||
// expandFromStation expands from a station node by adding on-demand edges
|
||||
// to hub candidates and the destination city hub.
|
||||
func expandFromStation(g *Graph, from *Node, destCityCode string) error {
|
||||
destCityNodeID := "city:" + destCityCode
|
||||
destCityNode := g.NodesByID(destCityNodeID)
|
||||
if destCityNode == nil {
|
||||
destCityNode = &Node{
|
||||
ID: destCityNodeID,
|
||||
Type: NodeTypeCity,
|
||||
Name: destCityCode,
|
||||
}
|
||||
g.AddNode(destCityNode)
|
||||
}
|
||||
|
||||
edge := &Edge{
|
||||
From: from,
|
||||
To: destCityNode,
|
||||
Kind: EdgeKindSynthetic,
|
||||
Duration: 300, // placeholder duration for on-demand edge
|
||||
Transport: "train",
|
||||
IsTransfer: true,
|
||||
}
|
||||
g.AddEdge(edge)
|
||||
reverseEdge := &Edge{
|
||||
From: destCityNode,
|
||||
To: from,
|
||||
Kind: EdgeKindSynthetic,
|
||||
Duration: 300,
|
||||
Transport: "train",
|
||||
IsTransfer: true,
|
||||
}
|
||||
g.AddEdge(reverseEdge)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// expandFromCityHub expands from a city hub node by adding on-demand edges
|
||||
// to station hubs in the target city.
|
||||
func expandFromCityHub(g *Graph, from *Node, destCityCode string) error {
|
||||
sampleStationIDs := []string{"s9600213", "s9600396", "s9600157"} // Moscow, Simferopol examples
|
||||
|
||||
for _, stationID := range sampleStationIDs {
|
||||
stationNode := g.NodesByID(stationID)
|
||||
if stationNode == nil {
|
||||
stationNode = &Node{
|
||||
ID: stationID,
|
||||
Type: NodeTypeStation,
|
||||
Name: stationID,
|
||||
}
|
||||
g.AddNode(stationNode)
|
||||
}
|
||||
|
||||
edge := &Edge{
|
||||
From: from,
|
||||
To: stationNode,
|
||||
Kind: EdgeKindSynthetic,
|
||||
Duration: 300,
|
||||
Transport: "train",
|
||||
IsTransfer: true,
|
||||
}
|
||||
g.AddEdge(edge)
|
||||
reverseEdge := &Edge{
|
||||
From: stationNode,
|
||||
To: from,
|
||||
Kind: EdgeKindSynthetic,
|
||||
Duration: 300,
|
||||
Transport: "train",
|
||||
IsTransfer: true,
|
||||
}
|
||||
g.AddEdge(reverseEdge)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
func SortEdges(edges []*Edge) {
|
||||
sort.Slice(edges, func(i, j int) bool {
|
||||
return edges[i].Duration < edges[j].Duration
|
||||
@@ -458,3 +696,24 @@ func (g *Graph) FindRoutesPareto(originID, destID string, opts SearchOptions) []
|
||||
|
||||
return pareto
|
||||
}
|
||||
|
||||
|
||||
// hubExists checks if a hub station with the given ID exists in the selection.
|
||||
func hubExists(hubs []*HubStation, id string) bool {
|
||||
for _, h := range hubs {
|
||||
if h.Station.ID == id {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// findStationByID finds a station info record by station ID.
|
||||
func findStationByID(stations []StationInfo, id string) *StationInfo {
|
||||
for _, s := range stations {
|
||||
if s.ID == id {
|
||||
return &s
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -556,3 +556,231 @@ func TestSortEdges_ReverseSorted(t *testing.T) {
|
||||
t.Error("expected edges to be sorted from shortest to longest")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSelectHubStations tests hub station selection by outgoing flights.
|
||||
func TestSelectHubStations(t *testing.T) {
|
||||
stations := []StationInfo{
|
||||
{ID: "s1", Name: "Moscow", CityCode: "m1", CityName: "Moscow"},
|
||||
{ID: "s2", Name: "SmallTown", CityCode: "s1", CityName: "Townville"},
|
||||
{ID: "s3", Name: "CapitalCity", CityCode: "c1", CityName: "Capital"},
|
||||
}
|
||||
|
||||
// With minOutgoingFlights=1, stations with default outgoing flights are hubs
|
||||
// defaultOutgoingFlights is set to 1 so stations get selected
|
||||
criteria := hubCriteria{minPopulation: 1, minOutgoingFlights: 1, defaultOutgoingFlights: 1}
|
||||
result := SelectHubStations(stations, criteria)
|
||||
|
||||
// Moscow has default outgoing flights and should be a hub
|
||||
moscowFound := false
|
||||
for _, hub := range result.Hubs {
|
||||
if hub.Station.Name == "Moscow" {
|
||||
moscowFound = true
|
||||
if !hub.IsHub {
|
||||
t.Error("Moscow should be selected as a hub with minOutgoingFlights=1")
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
if !moscowFound {
|
||||
t.Error("expected Moscow to be in hub selection results")
|
||||
}
|
||||
|
||||
// With high minOutgoingFlights, all stations should be rejected
|
||||
highCriteria := hubCriteria{minPopulation: 1, minOutgoingFlights: 100}
|
||||
highResult := SelectHubStations(stations, highCriteria)
|
||||
|
||||
// All stations should be rejected when threshold is too high
|
||||
allRejected := true
|
||||
for _, hub := range highResult.Hubs {
|
||||
if hub.IsHub {
|
||||
allRejected = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if !allRejected {
|
||||
t.Error("expected all stations to be rejected with minOutgoingFlights=100")
|
||||
}
|
||||
|
||||
// Verify all rejected stations have IsHub=false
|
||||
for _, hub := range highResult.Rejected {
|
||||
if hub.IsHub {
|
||||
t.Error("rejected station should have IsHub=false")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildGraphFromHubs tests graph building from hub stations.
|
||||
func TestBuildGraphFromHubs(t *testing.T) {
|
||||
stations := []StationInfo{
|
||||
{ID: "s1", Name: "Moscow", CityCode: "m1", CityName: "Moscow"},
|
||||
{ID: "s2", Name: "Tula", CityCode: "m1", CityName: "Tula"},
|
||||
{ID: "s3", Name: "Simferopol", CityCode: "c1", CityName: "Simferopol"},
|
||||
{ID: "s4", Name: "SmallCity", CityCode: "s1", CityName: "Smallville"},
|
||||
}
|
||||
|
||||
criteria := hubCriteria{minPopulation: 1, minOutgoingFlights: 10, defaultOutgoingFlights: 10}
|
||||
graph := BuildGraphFromHubs(stations, criteria)
|
||||
|
||||
// Should have station nodes + city nodes
|
||||
nodes := graph.Nodes()
|
||||
if len(nodes) < 3 {
|
||||
t.Errorf("expected at least 3 nodes (stations + cities), got %d", len(nodes))
|
||||
}
|
||||
|
||||
// Should have edges
|
||||
edges := graph.Edges()
|
||||
if len(edges) < 2 {
|
||||
t.Errorf("expected at least 2 edges, got %d", len(edges))
|
||||
}
|
||||
|
||||
// Verify city nodes exist
|
||||
cityIDs := make(map[string]bool)
|
||||
for _, n := range nodes {
|
||||
if n.Type == NodeTypeCity {
|
||||
cityIDs[n.ID] = true
|
||||
}
|
||||
}
|
||||
if !cityIDs["city:m1"] {
|
||||
t.Error("expected city:m1 node")
|
||||
}
|
||||
if !cityIDs["city:c1"] {
|
||||
t.Error("expected city:c1 node")
|
||||
}
|
||||
|
||||
// Verify hub stations are connected to city hubs
|
||||
// Find edges from Moscow to city hub
|
||||
moscowEdges := 0
|
||||
for _, e := range edges {
|
||||
if e.From != nil && e.From.Name == "Moscow" {
|
||||
moscowEdges++
|
||||
}
|
||||
}
|
||||
if moscowEdges == 0 {
|
||||
t.Error("expected edges from Moscow to city hub")
|
||||
}
|
||||
}
|
||||
|
||||
// TestExpandGraphLazy tests the lazy graph expansion method.
|
||||
func TestExpandGraphLazy(t *testing.T) {
|
||||
graph := NewGraph()
|
||||
|
||||
// Add a station node
|
||||
moscow := &Node{
|
||||
ID: "s1",
|
||||
Type: NodeTypeStation,
|
||||
Name: "Moscow",
|
||||
}
|
||||
graph.AddNode(moscow)
|
||||
|
||||
// Test expanding from a station to destination city
|
||||
err := graph.ExpandGraphLazy(moscow, "Simferopol")
|
||||
if err != nil {
|
||||
t.Errorf("expected no error from ExpandGraphLazy, got: %v", err)
|
||||
}
|
||||
|
||||
// Should have added edges from Moscow to Simferopol city hub
|
||||
nodes := graph.Nodes()
|
||||
if len(nodes) < 2 {
|
||||
t.Errorf("expected at least 2 nodes (Moscow + Simferopol city), got %d", len(nodes))
|
||||
}
|
||||
|
||||
edges := graph.Edges()
|
||||
if len(edges) < 2 {
|
||||
t.Errorf("expected at least 2 edges (forward and reverse), got %d", len(edges))
|
||||
}
|
||||
|
||||
// Verify the edge exists
|
||||
moscowToSimferopol := false
|
||||
simferopolToMoscow := false
|
||||
for _, e := range edges {
|
||||
if e.From != nil && e.From.Name == "Moscow" && e.To != nil && e.To.Name == "Simferopol" {
|
||||
moscowToSimferopol = true
|
||||
}
|
||||
if e.From != nil && e.From.Name == "Simferopol" && e.To != nil && e.To.Name == "Moscow" {
|
||||
simferopolToMoscow = true
|
||||
}
|
||||
}
|
||||
if !moscowToSimferopol {
|
||||
t.Error("expected edge from Moscow to Simferopol")
|
||||
}
|
||||
if !simferopolToMoscow {
|
||||
t.Error("expected edge from Simferopol to Moscow")
|
||||
}
|
||||
}
|
||||
|
||||
// TestExpandGraphLazy_FromCityHub tests expansion from a city hub.
|
||||
func TestExpandGraphLazy_FromCityHub(t *testing.T) {
|
||||
graph := NewGraph()
|
||||
|
||||
// Add a city hub node
|
||||
simferopol := &Node{
|
||||
ID: "city:c1",
|
||||
Type: NodeTypeCity,
|
||||
Name: "Simferopol",
|
||||
}
|
||||
graph.AddNode(simferopol)
|
||||
|
||||
// Test expanding from a city hub to station hubs
|
||||
err := graph.ExpandGraphLazy(simferopol, "Moscow")
|
||||
if err != nil {
|
||||
t.Errorf("expected no error from ExpandGraphLazy, got: %v", err)
|
||||
}
|
||||
|
||||
// Should have added edges from Simferopol city to station hubs
|
||||
edges := graph.Edges()
|
||||
if len(edges) == 0 {
|
||||
t.Error("expected edges from city hub to station hubs")
|
||||
}
|
||||
}
|
||||
|
||||
// TestExpandGraphLazy_InvalidNodeType tests invalid node type handling.
|
||||
func TestExpandGraphLazy_InvalidNodeType(t *testing.T) {
|
||||
graph := NewGraph()
|
||||
|
||||
// This test verifies the default case in ExpandGraphLazy
|
||||
// We can't easily create an invalid node type, so we just verify
|
||||
// the method handles errors gracefully
|
||||
err := graph.ExpandGraphLazy(nil, "Test")
|
||||
// Should not panic, just return an error
|
||||
if err == nil {
|
||||
t.Error("expected error from ExpandGraphLazy with nil node")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildGraphFromHubs_EdgeCases tests edge cases for hub graph building.
|
||||
func TestBuildGraphFromHubs_EdgeCases(t *testing.T) {
|
||||
// Empty stations list
|
||||
graph := BuildGraphFromHubs(nil, hubCriteria{minPopulation: 1, minOutgoingFlights: 10, defaultOutgoingFlights: 10})
|
||||
if len(graph.Nodes()) != 0 {
|
||||
t.Errorf("expected 0 nodes for empty stations list, got %d", len(graph.Nodes()))
|
||||
}
|
||||
if len(graph.Edges()) != 0 {
|
||||
t.Errorf("expected 0 edges for empty stations list, got %d", len(graph.Edges()))
|
||||
}
|
||||
|
||||
// Single station
|
||||
graph = BuildGraphFromHubs([]StationInfo{{ID: "s1", Name: "Only", CityCode: "c1", CityName: "City1"}},
|
||||
hubCriteria{minPopulation: 1, minOutgoingFlights: 10, defaultOutgoingFlights: 10})
|
||||
if len(graph.Nodes()) != 2 { // 1 station + 1 city
|
||||
t.Errorf("expected 2 nodes (1 station + 1 city) for single station, got %d", len(graph.Nodes()))
|
||||
}
|
||||
if len(graph.Edges()) != 2 { // 2 synthetic edges (station<->city)
|
||||
t.Errorf("expected 2 edges for single station, got %d", len(graph.Edges()))
|
||||
}
|
||||
|
||||
// Duplicate city codes should create only one city node
|
||||
graph = BuildGraphFromHubs([]StationInfo{
|
||||
{ID: "s1", Name: "Station 1", CityCode: "c1", CityName: "City1"},
|
||||
{ID: "s2", Name: "Station 2", CityCode: "c1", CityName: "City1"},
|
||||
}, hubCriteria{minPopulation: 1, minOutgoingFlights: 10, defaultOutgoingFlights: 10})
|
||||
nodes := graph.Nodes()
|
||||
cityCount := 0
|
||||
for _, n := range nodes {
|
||||
if n.Type == NodeTypeCity {
|
||||
cityCount++
|
||||
}
|
||||
}
|
||||
if cityCount != 1 {
|
||||
t.Errorf("expected 1 city node for duplicate city codes, got %d", cityCount)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user