Write nested loops as linq query

Question

I am tracking trains and trying to identify individual trains seen multiple times at different points through the IDs of the wagons on them when spotted.

// create a lookup of all tracked trains for each wagon
IEnumerable<Train> trains = GetTrains();

var wagonTrains = new Dictionary<int, List<Train>>();
foreach (Train t in trains)
{
  foreach (int w in t.WagonsInTrain)
  {
    if (!wagonTrains.ContainsKey(w))
    {
      wagonTrains.Add(w, new List<Train>());
    }
    wagonTrains[w].Add(t);
  }
}

Is there a better way do to what I am doing in this code segment? Perhaps some chained linq operations? Is there a particular name for the kind of operation I am using here?

Answer 1

create a lookup of all tracked trains for each wagon

That's pretty much what the ToLookup() method is for. You just need a bit more LINQ to get a collection of (train, wagon) pairs, so that ToLookup() can work:

var wagonTrains =
    (from train in trains
     from wagon in train.WagonsInTrain
     select new { train, wagon })
        .ToLookup(x => x.wagon, x => x.train);

Answer 2

Here is one that uses chained linq operations:

var trainsByWagon =
    trains
        .SelectMany(train => train.WagonsInTrain, (train, wagon) => new { train, wagon })
        .GroupBy(trainAndWagon => trainAndWagon.wagon, trainAndWagon => trainAndWagon.train);
        .ToDictionary(g => g.Key, g => g.ToList());

Small print: Haven't actually tried this, but should work.

Answer 3

I've actually set this up in my code as a pair of extension methods:

public static class IDictionaryExtensions
{
    public static void Update<TKEY, TVALUE>(this IDictionary<TKEY, TVALUE> dictionary, TKEY key, TVALUE value)
    {
        if (dictionary.ContainsKey(key))
        {
            dictionary[key] = value;
        }
        else
        {
            dictionary.Add(key, value);
        }
    }
    public static void AddToList<TKEY, TVALUE>(this IDictionary<TKEY, IList<TVALUE>> dictionary, TKEY key, TVALUE value)
    {
        if (dictionary.ContainsKey(key) && dictionary[key] != null)
        {
            dictionary[key].Add(value);
        }
        else
        {
            dictionary.Update(key, new List<TVALUE> {value});
        }
    }
 }

The first method is an all-purpose "Add"-type method which will either add it or update it depending on whether or not it exists. The second is specifically for dictionaries where the value is a List<>. It'll let you add an element to the List<>, creating the key if neccesary.

---

For your example, usage would be:

IEnumerable<Train> trains = GetTrains();
var wagonTrains = new Dictionary<int, List<Train>>();
foreach (var train in trains)
{
   foreach (var wagon in train.WagonsInTrain)
   {
       wagonTrains.AddToList(wagon, train);
   }
}